Merge branch 'master' into remove_vec_2

.github/workflows/benchmark.yml

@@ -623,6 +623,8 @@ jobs:
       run: test/external/process_replay/reset.py
     - name: openpilot compile3 0.11.0 driving_vision
       run: BENCHMARK_LOG=openpilot_0_11_0_vision PYTHONPATH="." ASSERT_MIN_STEP_TIME=17 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.11.0/selfdrive/modeld/models/driving_vision.onnx
+    - name: openpilot compile3 0.11.0 driving_vision (from pickle)
+      run: BENCHMARK_LOG=openpilot_0_11_0_vision_run_pickle RUN_PICKLE=1 PYTHONPATH="." ASSERT_MIN_STEP_TIME=17 DEV=QCOM taskset -c 4-7 python3 examples/openpilot/compile3.py
     - name: IR3 openpilot compile3 0.11.0 driving_vision
       run: BENCHMARK_LOG=ir3_openpilot_0_11_0_vision PYTHONPATH="." ASSERT_MIN_STEP_TIME=17 DEV=QCOM:IR3 FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.11.0/selfdrive/modeld/models/driving_vision.onnx
     - name: openpilot compile3 0.11.0 driving_policy
@@ -668,6 +670,8 @@ jobs:
       run: BENCHMARK_LOG=usbgpu_openpilot_0_10_1_vision PYTHONPATH="." GMMU=0 DEV=USB+AMD:LLVM ASSERT_MIN_STEP_TIME=50 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
     - name: openpilot load_pickle 0.10.1 driving_vision
       run: BENCHMARK_LOG=usbgpu_openpilot_0_10_1_vision_load_pickle PYTHONPATH="." GMMU=0 DEV=USB+AMD ASSERT_MIN_LOAD_TIME=15 python3 examples/openpilot/load_pickle.py
+    - name: openpilot run_pickle 0.10.1 driving_vision
+      run: BENCHMARK_LOG=usbgpu_openpilot_0_10_1_vision_run_pickle RUN_PICKLE=1 PYTHONPATH="." GMMU=0 DEV=USB+AMD ASSERT_MIN_STEP_TIME=50 python3 examples/openpilot/compile3.py
 
   testreddriverbenchmark:
     name: AM Benchmark

.gitignore

@@ -68,3 +68,4 @@ mutants
 .mutmut-cache
 dagre/
 graphlib/
+uv.lock

docs/abstractions3.py

@@ -1,6 +1,4 @@
 # abstractions2 goes from back to front, here we will go from front to back
-from typing import List
-from tinygrad.helpers import tqdm
 
 # *****
 # 0. Load mnist on the device
@@ -33,21 +31,21 @@ model(X).sparse_categorical_crossentropy(Y).backward()
 optim.schedule_step()   # this will step the optimizer without running realize
 
 # *****
-# 3. Create a schedule.
+# 3. Create a schedule (linear uop).
 
 # The weight Tensors have been assigned to, but not yet realized. Everything is still lazy at this point
 # l1.uop and l2.uop define a computation graph
 
-from tinygrad.schedule import ExecItem
-schedule: List[ExecItem] = Tensor.schedule(l1, l2)
+from tinygrad.engine.realize import run_linear
+linear = Tensor.schedule_linear(l1, l2)
 
-print(f"The schedule contains {len(schedule)} items.")
-for si in schedule: print(str(si)[:80])
+print(f"The schedule contains {len(linear.src)} items.")
+for call in linear.src: print(str(call)[:80])
 
 # *****
-# 4. Lower and run the schedule.
+# 4. Lower and run the schedule (linear uop).
 
-for si in tqdm(schedule): si.run()
+run_linear(linear)
 
 # *****
 # 5. Print the weight change

docs/developer/developer.md

@@ -17,13 +17,11 @@ The `UOp` graph specifies the compute in terms of low level tinygrad ops. Not al
 
 ## Scheduling
 
-The [scheduler](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/schedule/__init__.py) converts the graph of UOps into a list of `ExecItem`. One `ExecItem` is one kernel on the GPU, and the scheduler is responsible for breaking the large compute graph into subgraphs that can fit in a kernel. `ast` specifies what compute to run, and `bufs` specifies what buffers to run it on.
-
-::: tinygrad.schedule.ExecItem
+The [scheduler](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/schedule/__init__.py) converts the graph of UOps into a `LINEAR` UOp whose `src` is a list of `CALL` UOps. One `CALL` is one kernel on the GPU, and the scheduler is responsible for breaking the large compute graph into subgraphs that can fit in a kernel. The `CALL`'s `src[0]` (a `SINK` ast) specifies what compute to run, and the remaining `src` are the buffers to run it on.
 
 ## Lowering
 
-The code in [realize](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/engine/realize.py) lowers `ExecItem` by populating its `prg` field with
+The code in [realize](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/engine/realize.py) lowers each `CALL` by compiling its ast into a `PROGRAM` and running it.
 
 ::: tinygrad.engine.realize.run_linear
 
@@ -35,13 +33,7 @@ Then we render the UOps into code with a `Renderer`, then we compile the code to
 
 ## Execution
 
-Creating `ExecItem`, which has a run method
-
-::: tinygrad.engine.realize.ExecItem
-    options:
-        members: true
-
-Lists of `ExecItem` can be condensed into a single ExecItem with the Graph API (rename to Queue?)
+`run_linear` walks the `LINEAR` UOp, dispatching each `CALL` to a runner (kernel, copy, view, encdec, or graph).
 
 ## Runtime
 

docs/developer/layout.md

@@ -28,7 +28,7 @@ Transforms the ast into an optimized ast. This is where BEAM search and heuristi
 
 Transform the optimized ast into a linearized and rendered program.
 
-::: tinygrad.codegen.get_program
+::: tinygrad.codegen.to_program
     options:
         members: false
         show_labels: false
@@ -53,7 +53,7 @@ Transform the linearized list of UOps into a program, represented as a string.
 
 Abstracted high level interface to the runtimes.
 
-::: tinygrad.engine.realize.get_program
+::: tinygrad.engine.realize.to_program
     options:
         members: false
         show_labels: false

docs/tensor/properties.md

@@ -19,8 +19,8 @@
 
 ## tinygrad ops
 
-::: tinygrad.Tensor.schedule_with_vars
-::: tinygrad.Tensor.schedule
+::: tinygrad.Tensor.linear_with_vars
+::: tinygrad.Tensor.schedule_linear
 ::: tinygrad.Tensor.realize
 ::: tinygrad.Tensor.replace
 ::: tinygrad.Tensor.assign

examples/anthropic_challenge.py

@@ -173,16 +173,16 @@ if __name__ == "__main__":
 
   # *** render to device ***
 
-  from tinygrad.codegen import get_program
+  from tinygrad.codegen import to_program
   with Context(PCONTIG=2, DEVECTORIZE=2, SPEC=0):
     out = tree_traversal(forest_t, val_t, height, rounds)
-    sink = out.schedule()[-1].ast
-    prg = get_program(sink, VLIWRenderer())
+    sink = out.schedule_linear().src[-1].src[0]
+    prg = to_program(sink, VLIWRenderer())
 
   # *** run on Machine and compare ***
 
   # NOTE: the scratch size needs to be reduced to 1536 when you have a register allocator
-  src = eval(prg.src)
+  src = eval(prg.src[3].arg)
   max_regs = max(t[1] for instr in src for v in instr.values() for t in v if len(t) > 1) + 8
   print(f"{max_regs:5d} regs used" + ("" if max_regs <= 1536 else "       <-- WARNING: TOO MANY REGISTERS, MUST BE <= 1536"))
   machine = problem.Machine(mem, src, problem.DebugInfo(scratch_map={}), n_cores=1, trace=False, scratch_size=max_regs)

examples/mlperf/model_train.py

@@ -1282,7 +1282,7 @@ def train_bert():
         previous_step = i
 
 def train_llama3():
-  from examples.mlperf.models.flat_llama import FlatTransformer, apply_grad, FP8, FP8_DTYPE
+  from examples.mlperf.models.flat_llama import FlatTransformer, apply_grad, FP8_DTYPE
   from examples.llama3 import MODEL_PARAMS
   from examples.mlperf.lr_schedulers import CosineAnnealingLRWithWarmup
   from examples.mlperf.optim import GradAccClipAdamW
@@ -1395,7 +1395,7 @@ def train_llama3():
 
   params = get_parameters(model)
 
-  if getenv("FAKEDATA"):
+  if getenv("EMPTYWEIGHT"):
     for v in get_parameters(model):
       v = v.assign(Tensor.empty(v.shape, dtype=v.dtype))
 
@@ -1432,18 +1432,17 @@ def train_llama3():
     print(f"loading optim checkpoint from {fn}")
     load_state_dict(scheduler, safe_load(fn), realize=False)
 
-  fp8_amax = [t for ts in model._fp8_amax.values() for t in ts] if FP8 else []
-  fp8_inv_scales = list(model._fp8_inv_scale.values()) if FP8 else []
+  fp8_amax = [t for ts in model._fp8_amax.values() for t in ts]
+  fp8_inv_scales = list(model._fp8_inv_scale.values())
 
-  if FP8:
-    from tinygrad.nn.state import get_state_dict
-    model_state = get_state_dict(model)
-    for wname in ["wqkv", "wo", "w13", "w2"]:
-      w = model_state[wname]
-      w._inv_scale = model._fp8_inv_scale[wname]
-      if optim.master_params:
-        idx = next(j for j, p in enumerate(optim.params) if p is w)
-        optim.master_params[idx].assign((optim.master_params[idx] * w._inv_scale.reshape(-1, *([1]*(w.ndim-1)))).contiguous())
+  from tinygrad.nn.state import get_state_dict
+  model_state = get_state_dict(model)
+  for wname in ["wqkv", "wo", "w13", "w2"]:
+    w = model_state[wname]
+    w._inv_scale = model._fp8_inv_scale[wname]
+    if optim.master_params:
+      idx = next(j for j, p in enumerate(optim.params) if p is w)
+      optim.master_params[idx].assign((optim.master_params[idx] * w._inv_scale.reshape(-1, *([1]*(w.ndim-1)))).contiguous())
 
   @TinyJit
   def minibatch(tokens:Tensor):
@@ -1451,7 +1450,11 @@ def train_llama3():
     if is_mp: tokens = tokens.shard(device)
     if not is_sharding: tokens = tokens.to(None)
     logits:Tensor = model(tokens[:, :-1])
-    loss = vocab_mask.where(-1e9, logits).sparse_categorical_crossentropy(tokens[:, 1:])
+    if getenv("FAST_CE", 0):
+      from extra.llama_kernels.fused_ce import fused_ce_loss
+      loss = fused_ce_loss(logits.cast(dtypes.bfloat16), tokens[:, 1:], label_smoothing=0.0)
+    else:
+      loss = vocab_mask.where(-1e9, logits).sparse_categorical_crossentropy(tokens[:, 1:])
 
     for g, new_g in zip(grads, loss.gradient(*optim.params)):
       apply_grad(g, new_g.uop)
@@ -1555,7 +1558,7 @@ def train_llama3():
 
       mem_gb = GlobalCounters.mem_used / 1e9
       gflops = GlobalCounters.global_ops / 1e9 / dev_time
-      mfu = ((6 * num_params * SEQLEN * GBS) / (dev_time * device_count * (4.6e15 if FP8 else 2.3e15))) * 100
+      mfu = ((6 * num_params * SEQLEN * GBS) / (dev_time * device_count * 4.6e15)) * 100
       tqdm.write(
           f"{i:5} {step_time:.3f} s step, {gbs_time:.3f} s gbs, {optim_time:.3f} s optim, {data_time:.3f} s data, {loss:.4f} loss, " \
           f"{lr:.12f} LR, {grad_norm:.6f} grad_norm, {mem_gb:.2f} GB used, {gflops:9.2f} GFLOPS, {mfu:5.2f}% MFU")

examples/mlperf/models/flat_llama.py

@@ -1,4 +1,4 @@
-import math, os, functools
+import math, os
 if __name__ == "__main__":
   os.environ["DEFAULT_FLOAT"] = "bfloat16"
   os.environ["OPTIM_DTYPE"] = "bfloat16"
@@ -16,65 +16,60 @@ from tinygrad import Tensor, nn, function, getenv, dtypes, TinyJit
 from tinygrad.helpers import Timing, colored, GlobalCounters, profile_marker
 from tinygrad.uop.ops import Ops, UOp
 from extra.models.llama import apply_rotary_emb, precompute_freqs_cis
+from extra.llama_kernels.rmsnorm import rmsnorm
+from extra.llama_kernels import FP8_MAX, local_abs_max
 
-FP8 = getenv("FP8", 0)
+ASM_GEMM = getenv("ASM_GEMM", 0)
 
 FP8_DTYPE = dtypes.fp8e4m3
 FP8_GRAD_DTYPE = dtypes.fp8e5m2
-FP8_MAX = 448.0
-
-# per-device abs max without allreduce (matches TE delayed scaling behavior)
-@functools.cache
-def _local_abs_max_fxn(x_p, device):
-  x = Tensor(x_p, device=device)
-  inner = Tensor(x.uop.src[0]) if x.uop.op is Ops.MULTI else x
-  return (inner.abs().max(),)
-
-def _local_abs_max(x:Tensor) -> Tensor:
-  param = x.as_param(0)
-  fxn = _local_abs_max_fxn(param.uop, x.device)
-  return Tensor(fxn[0].uop.call(x.uop).gettuple(0))
 
 def quantize_fp8(x:Tensor, amax_state:Tensor|None=None):
-  new_amax = (_local_abs_max(x) if isinstance(x.device, tuple) else x.abs().max()).detach()
+  new_amax = (local_abs_max(x) if isinstance(x.device, tuple) else x.abs().max()).detach()
   scale = FP8_MAX / ((amax_state if amax_state is not None else new_amax) + 1e-8)
   x_scaled = x * scale
   x_clamped = x_scaled + (x_scaled.detach().clamp(-FP8_MAX, FP8_MAX) - x_scaled.detach())  # STE
   return x_clamped.cast(FP8_DTYPE), scale.float().reciprocal(), new_amax
 
-def matmul(x:Tensor, w:Tensor, fp8=FP8, amax_x:Tensor|None=None, w_inv_scale:Tensor|None=None,
+def matmul(x:Tensor, w:Tensor, fp8:bool=True, amax_x:Tensor|None=None, w_inv_scale:Tensor|None=None,
            x_fp8:Tensor|None=None, x_scale:Tensor|None=None, x_new_amax:Tensor|None=None) -> tuple[Tensor,...]:
   if not fp8:
-    if getenv("ASM_GEMM"):
+    if ASM_GEMM:
       from extra.gemm.cdna_asm_gemm import can_use_asm_gemm, asm_gemm
       if can_use_asm_gemm(x, w.T): return (asm_gemm(x, w.T),)
     return (x @ w.T,)
   assert w_inv_scale is not None, "fp8 matmul requires w_inv_scale (weights must be stored in fp8 with per-tensor scale)"
   if x_fp8 is None: x_fp8, x_scale, x_new_amax = quantize_fp8(x, amax_state=amax_x)
-  if getenv("ASM_GEMM"):
+  if ASM_GEMM:
     from extra.gemm.cdna_asm_gemm import can_use_asm_gemm, asm_gemm
     if can_use_asm_gemm(x_fp8, w.T): return asm_gemm(x_fp8, w.T, x_scale=x_scale, w_scale=w_inv_scale), x_new_amax, x_fp8, w
-  return x_fp8.dot(w.T, dtype=dtypes.float) * x_scale * w_inv_scale, x_new_amax, x_fp8, w
+  return (x_fp8.dot(w.T, dtype=dtypes.float) * x_scale * w_inv_scale).cast(dtypes.bfloat16), x_new_amax, x_fp8, w
 
-def _rmsnorm_fwd(x_in:Tensor, eps:float) -> tuple[Tensor, Tensor]:
-  x = x_in.float()
-  rrms = (x.square().mean(-1, keepdim=True) + eps).rsqrt()
-  return (x * rrms).cast(x_in.dtype), rrms
+def norm_mul_quantize_matmul(x:Tensor, norm:Tensor, amax_x, w_inv_scale, w:Tensor, eps:float):
+  FUSED_NORM_MUL_QUANTIZE = getenv("FUSED_NORM_MUL_QUANTIZE", 0)
+  normed, rrms = rmsnorm(x, eps)
+  if FUSED_NORM_MUL_QUANTIZE:
+    from extra.llama_kernels.fused_mul_quantize_fp8 import fused_mul_quantize_fp8
+    amax_s = amax_x if amax_x is not None else Tensor.full((), 1.0, dtype=dtypes.bfloat16, device=normed.device)
+    x_fp8, x_inv_scale, new_amax = fused_mul_quantize_fp8(normed, norm, amax_s, FP8_DTYPE)
+    out, *ret = matmul(None, w, w_inv_scale=w_inv_scale, x_fp8=x_fp8, x_scale=x_inv_scale, x_new_amax=new_amax)
+  else:
+    x = normed * norm
+    out, *ret = matmul(x, w, amax_x=amax_x, w_inv_scale=w_inv_scale)
+  return out, normed, rrms, ret
 
-@functools.cache
-def _rmsnorm_fwd_fxn(x_in_p, eps, device):
-  return _rmsnorm_fwd(Tensor(x_in_p, device=device), eps)
-
-def _rmsnorm_bwd(grad:UOp, call:UOp) -> tuple:
-  x_normed = Tensor(call.gettuple(0)).float()
-  do_float = Tensor(grad).float()
-  d_x = Tensor(call.gettuple(1)) * (do_float - x_normed * (do_float * x_normed).mean(-1, keepdim=True))
-  return (d_x.cast(call.src[1].dtype).uop,)
-
-def rmsnorm(x_in:Tensor, eps:float) -> tuple[Tensor, Tensor]:
-  fxn = _rmsnorm_fwd_fxn(x_in.as_param(0).uop, eps, x_in.device)
-  call = UOp.maketuple(fxn[0].uop, fxn[1].uop).call(x_in.uop, grad_fxn=_rmsnorm_bwd)
-  return Tensor(call.gettuple(0)), Tensor(call.gettuple(1))
+def silu_w13_matmul(x_w13:Tensor, w2:Tensor, amax_x2, s_2):
+  FUSED_SILU_W13 = getenv("FUSED_SILU_W13", 0)
+  if FUSED_SILU_W13:
+    from extra.llama_kernels.cast_amax import fused_quantize_fp8_w13
+    amax_s = amax_x2 if amax_x2 is not None else Tensor.full((), 1.0, dtype=dtypes.bfloat16, device=x_w13.device)
+    x2_fp8, x2_inv_scale, new_amax_x2 = fused_quantize_fp8_w13(x_w13, amax_s, FP8_DTYPE)
+    out, *ret = matmul(None, w2, w_inv_scale=s_2, x_fp8=x2_fp8, x_scale=x2_inv_scale, x_new_amax=new_amax_x2)
+  else:
+    hidden_dim = x_w13.shape[-1] // 2
+    x_w1, x_w3 = x_w13[..., :hidden_dim], x_w13[..., hidden_dim:]
+    out, *ret = matmul(x_w1.silu() * x_w3, w2, amax_x=amax_x2, w_inv_scale=s_2)
+  return out, ret
 
 class FlatTransformer:
   def __init__(self, dim:int, hidden_dim:int, n_heads:int, n_layers:int, norm_eps:float, vocab_size:int, n_kv_heads:int|None=None,
@@ -90,7 +85,7 @@ class FlatTransformer:
     scaled_std = 0.02 / math.sqrt(2 * n_layers)
 
     # Attention
-    self._init_inv_scales = []  # populated by lin_per_layer when FP8
+    self._init_inv_scales = []  # populated by lin_per_layer
     self.wqkv = self.lin_per_layer(dim, self.n_heads * self.head_dim + self.n_kv_heads * self.head_dim * 2)
     self.wo = self.lin_per_layer(self.n_heads * self.head_dim, dim, std=scaled_std)
 
@@ -109,21 +104,19 @@ class FlatTransformer:
     self.output = Tensor.normal(1, vocab_size, dim, mean=0.0, std=0.02, dtype=dtypes.bfloat16)
     self.freqs_cis = precompute_freqs_cis(dim // n_heads, max_context * 2, rope_theta).contiguous().requires_grad_(False)
 
-    if FP8:
-      def _amax(): return Tensor.full((), FP8_MAX).contiguous().requires_grad_(False)
-      names = ["xqkv", "xo", "x13", "x2"]
-      self._fp8_amax = {name: [_amax() for _ in range(n_layers)] for name in names}
-      # per-weight inv_scale: single (n_layers,) float32 tensor per weight (kernel reads float* pointers)
-      w_names = ["wqkv", "wo", "w13", "w2"]
-      self._fp8_inv_scale = {}
-      for wname, inv_scales in zip(w_names, self._init_inv_scales):
-        self._fp8_inv_scale[wname] = inv_scales.float().contiguous().requires_grad_(False)
-      del self._init_inv_scales
+    def _amax(): return Tensor.full((), FP8_MAX).contiguous().requires_grad_(False)
+    names = ["xqkv", "xo", "x13", "x2"]
+    self._fp8_amax = {name: [_amax() for _ in range(n_layers)] for name in names}
+    # per-weight inv_scale: single (n_layers,) float32 tensor per weight (kernel reads float* pointers)
+    w_names = ["wqkv", "wo", "w13", "w2"]
+    self._fp8_inv_scale = {}
+    for wname, inv_scales in zip(w_names, self._init_inv_scales):
+      self._fp8_inv_scale[wname] = inv_scales.float().contiguous().requires_grad_(False)
+    del self._init_inv_scales
 
   def lin_per_layer(self, in_features:int, out_features:int, std:float=0.02):
-    if getenv("ZEROS"): w = Tensor.zeros(self.n_layers, out_features, in_features)
+    if getenv("ZEROS", 0): w = Tensor.zeros(self.n_layers, out_features, in_features)
     else: w = Tensor.normal(self.n_layers, out_features, in_features, mean=0.0, std=std)
-    if not FP8: return w
     # per-layer scaled fp8 cast: fill the fp8 range for best precision
     amax = w.abs().flatten(1).max(1).detach()
     scale = FP8_MAX / (amax + 1e-8)
@@ -135,18 +128,8 @@ class FlatTransformer:
     bsz, seqlen, _ = x.shape
     new_amaxs, saves = [], []
 
-    x, rrms = rmsnorm(x, self.norm_eps)
-    saves.extend([x, rrms])
-
-    if FP8 and getenv("FUSED_NORM_MUL_QUANTIZE", 1):
-      from extra.amax.cast_amax import fused_mul_quantize_fp8
-      amax_s = amax_xqkv if amax_xqkv is not None else Tensor.full((), 1.0, dtype=dtypes.bfloat16, device=x.device)
-      x_fp8, x_inv_scale, new_amax_xqkv = fused_mul_quantize_fp8(x, attention_norm, amax_s, FP8_DTYPE)
-      xqkv, *ret = matmul(None, wqkv, w_inv_scale=s_qkv, x_fp8=x_fp8, x_scale=x_inv_scale, x_new_amax=new_amax_xqkv)
-    else:
-      x = x * attention_norm
-      xqkv, *ret = matmul(x, wqkv, amax_x=amax_xqkv, w_inv_scale=s_qkv)
-
+    xqkv, normed, rrms, ret = norm_mul_quantize_matmul(x, attention_norm, amax_xqkv, s_qkv, wqkv, self.norm_eps)
+    saves.extend([normed, rrms])
     new_amaxs.extend(ret[:1])
     saves.extend(ret[1:] + [xqkv])
     xqkv = xqkv.reshape(bsz, seqlen, self.n_kv_heads, self.n_rep + 2, self.head_dim)
@@ -155,7 +138,7 @@ class FlatTransformer:
     xv = xqkv[:, :, :, self.n_rep+1].reshape(bsz, seqlen, self.n_kv_heads, self.head_dim)
 
     xq, xk = apply_rotary_emb(xq, xk, freqs_cis)
-    if FP8: xq, xk, xv = xq.cast(dtypes.bfloat16), xk.cast(dtypes.bfloat16), xv.cast(dtypes.bfloat16)
+    xq, xk, xv = xq.cast(dtypes.bfloat16), xk.cast(dtypes.bfloat16), xv.cast(dtypes.bfloat16)
     xq, xk, xv = xq.transpose(1, 2), xk.transpose(1, 2), xv.transpose(1, 2)
     if getenv("HK_FLASH_ATTENTION"):
       from extra.thunder.amd.fa import flash_attention
@@ -174,28 +157,12 @@ class FlatTransformer:
                    amax_x13=None, amax_x2=None, s_13=None, s_2=None):
     new_amaxs, saves = [], []
 
-    x, rrms = rmsnorm(x, self.norm_eps)
-    saves.extend([x, rrms])
-
-    if FP8 and getenv("FUSED_NORM_MUL_QUANTIZE", 1):
-      from extra.amax.cast_amax import fused_mul_quantize_fp8
-      amax_s13 = amax_x13 if amax_x13 is not None else Tensor.full((), 1.0, dtype=dtypes.bfloat16, device=x.device)
-      x_fp8_13, x_inv_scale_13, new_amax_x13 = fused_mul_quantize_fp8(x, ffn_norm, amax_s13, FP8_DTYPE)
-      x_w13, *ret = matmul(None, w13, w_inv_scale=s_13, x_fp8=x_fp8_13, x_scale=x_inv_scale_13, x_new_amax=new_amax_x13)
-    else:
-      x = x * ffn_norm
-      x_w13, *ret = matmul(x, w13, amax_x=amax_x13, w_inv_scale=s_13)
+    x_w13, normed, rrms, ret = norm_mul_quantize_matmul(x, ffn_norm, amax_x13, s_13, w13, self.norm_eps)
+    saves.extend([normed, rrms])
     new_amaxs.extend(ret[:1])
     saves.extend(ret[1:] + [x_w13])
 
-    if FP8 and getenv("FUSED_SILU_W13", 1):
-      from extra.amax.cast_amax import fused_quantize_fp8_w13
-      amax_s = amax_x2 if amax_x2 is not None else Tensor.full((), 1.0, dtype=dtypes.bfloat16, device=x_w13.device)
-      x2_fp8, x2_inv_scale, new_amax_x2 = fused_quantize_fp8_w13(x_w13, amax_s, FP8_DTYPE)
-      out, *ret = matmul(None, w2, w_inv_scale=s_2, x_fp8=x2_fp8, x_scale=x2_inv_scale, x_new_amax=new_amax_x2)
-    else:
-      x_w1, x_w3 = x_w13[..., :self.hidden_dim], x_w13[..., self.hidden_dim:]
-      out, *ret = matmul(x_w1.silu() * x_w3, w2, amax_x=amax_x2, w_inv_scale=s_2)
+    out, ret = silu_w13_matmul(x_w13, w2, amax_x2, s_2)
     new_amaxs.extend(ret[:1])
     saves.extend(ret[1:] + [out])
     return (out, *new_amaxs, *saves)
@@ -226,41 +193,35 @@ class FlatTransformer:
     else:
       # flat per-layer weights: axis 0 is n_layers, so shard axes are +1 vs per-layer Transformer
       self.wqkv.shard_(device, axis=1).realize()          # (n_layers, out, dim) shard out
-      self.wo.shard_(device, axis=2).realize()             # (n_layers, dim, in) shard in
-      self.w13.shard_(device, axis=1).realize()             # (n_layers, hidden*2, dim) shard out
-      self.w2.shard_(device, axis=2).realize()             # (n_layers, dim, hidden) shard in
+      self.wo.shard_(device, axis=2).realize()            # (n_layers, dim, in) shard in
+      self.w13.shard_(device, axis=1).realize()           # (n_layers, hidden*2, dim) shard out
+      self.w2.shard_(device, axis=2).realize()            # (n_layers, dim, hidden) shard in
       self.attention_norm.shard_(device, axis=None).realize()
       self.ffn_norm.shard_(device, axis=None).realize()
       self.norm.weight.shard_(device, axis=None).realize()
       self.tok_embeddings.weight.shard_(device, axis=0).realize()
       self.output.shard_(device, axis=1).realize()
       self.freqs_cis.shard_(device, axis=None).realize()
-      if FP8:
-        for name in self._fp8_amax:
-          for i in range(len(self._fp8_amax[name])):
-            self._fp8_amax[name][i] = self._fp8_amax[name][i].to(device).contiguous().requires_grad_(False)
-        for name in self._fp8_inv_scale:
-          self._fp8_inv_scale[name] = self._fp8_inv_scale[name].to(device).contiguous().requires_grad_(False)
+      for name in self._fp8_amax:
+        for i in range(len(self._fp8_amax[name])):
+          self._fp8_amax[name][i] = self._fp8_amax[name][i].to(device).contiguous().requires_grad_(False)
+      for name in self._fp8_inv_scale:
+        self._fp8_inv_scale[name] = self._fp8_inv_scale[name].to(device).contiguous().requires_grad_(False)
 
   def __call__(self, tokens:Tensor):
     h = self.tok_embeddings(tokens)
     freqs_cis = self.freqs_cis.cast(h.dtype)[:, :tokens.shape[1], :, :, :]
-    a = self._fp8_amax if FP8 else None
-    s = self._fp8_inv_scale if FP8 else None
+    amaxs, inv_scales = self._fp8_amax, self._fp8_inv_scale
     for i in range(self.n_layers):
-      amax_layer = {"amax_xqkv": a["xqkv"][i], "amax_xo": a["xo"][i],
-                    "amax_x13": a["x13"][i], "amax_x2": a["x2"][i]} if a else {}
-      scale_layer = {"s_qkv": s["wqkv"][i], "s_o": s["wo"][i],
-                     "s_13": s["w13"][i], "s_2": s["w2"][i]} if s else {}
       h, *ret = self.run_layer(h, freqs_cis,
                                self.attention_norm[i], self.wqkv[i], self.wo[i],
                                self.ffn_norm[i], self.w13[i], self.w2[i],
-                               **amax_layer, **scale_layer)
-      if a:
-        amaxs = ret[:5]
-        amax_names = ["xqkv", "xo", "x13", "x2"]
-        for name, new_val in zip(amax_names, amaxs):
-          a[name][i].assign(new_val)
+                               amax_xqkv=amaxs["xqkv"][i], amax_xo=amaxs["xo"][i],
+                               amax_x13=amaxs["x13"][i], amax_x2=amaxs["x2"][i],
+                               s_qkv=inv_scales["wqkv"][i], s_o=inv_scales["wo"][i],
+                               s_13=inv_scales["w13"][i], s_2=inv_scales["w2"][i])
+      for name, new_val in zip(["xqkv", "xo", "x13", "x2"], ret[:5]):
+        amaxs[name][i].assign(new_val)
 
     logits = matmul(self.norm(h).contiguous().contiguous_backward(), self.output[0], fp8=False)[0].contiguous_backward()
     return logits

examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/dev_beam.sh

@@ -15,9 +15,10 @@ export WQKV=${WQKV:-1}
 export MASTER_WEIGHTS=${MASTER_WEIGHTS:-1}
 export FP8=${FP8:-1}
 export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1}
+export FAST_CE=${FASE_CE:-1}
 
 export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
-export DP=${DP:-8} MP=${MP:-1} BS=${BS:-16} EVAL_BS=${EVAL_BS:-16} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-2}
+export DP=${DP:-8} MP=${MP:-1} BS=${BS:-16} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-2}
 export GBS=$((BS * GRADIENT_ACC_STEPS))
 
 export MODEL="llama3"
@@ -36,7 +37,7 @@ export DATA_SEED=${DATA_SEED:-5760}
 export JITBEAM=${JITBEAM:-3}
 export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=1
 
-export FAKEDATA=1 BENCHMARK=${BENCHMARK:-10}
+export FAKEDATA=${FAKEDATA:-1} BENCHMARK=${BENCHMARK:-10}
 if [ -z "$FULL_LAYERS" ]; then
   export LLAMA_LAYERS=2
 fi

examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/dev_run.sh

@@ -15,9 +15,10 @@ export WQKV=${WQKV:-1}
 export MASTER_WEIGHTS=${MASTER_WEIGHTS:-1}
 export FP8=${FP8:-1}
 export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1}
+export FAST_CE=${FASE_CE:-1}
 
 export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
-export DP=${DP:-8} MP=${MP:-1} BS=${BS:-16} EVAL_BS=${EVAL_BS:-16} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-2}
+export DP=${DP:-8} MP=${MP:-1} BS=${BS:-16} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-2}
 export GBS=$((BS * GRADIENT_ACC_STEPS))
 
 export MODEL="llama3"

examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/profile.sh

@@ -3,4 +3,4 @@ export BENCHMARK=5
 export EVAL_BS=0
 VIZ=${VIZ:--1} FULL_LAYERS=1 DEBUG=0 examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/dev_beam.sh
 SRC="AMD"; [[ $DEV == NULL* ]] && SRC="NULL"
-python -m tinygrad.viz.cli --profile -s "$SRC" --top 20
+python -m tinygrad.viz.cli -s "$SRC" --top 20

examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/run_and_time.sh

@@ -16,9 +16,10 @@ export WQKV=1
 export MASTER_WEIGHTS=1
 export FP8=1
 export ALLREDUCE_CAST=1
+export FAST_CE=1
 
 export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
-export DP=8 MP=1 BS=16 EVAL_BS=16 GRADIENT_ACC_STEPS=2
+export DP=8 MP=1 BS=16 EVAL_BS=8 GRADIENT_ACC_STEPS=2
 export GBS=$((BS * GRADIENT_ACC_STEPS))
 
 export MODEL="llama3"

examples/openpilot/compile3.py

@@ -5,7 +5,6 @@ if "JIT_BATCH_SIZE" not in os.environ: os.environ["JIT_BATCH_SIZE"] = "0"
 from tinygrad import fetch, Tensor, TinyJit, Context, GlobalCounters, Device, dtypes
 from tinygrad.helpers import DEBUG, getenv
 from tinygrad.uop.ops import Ops
-from tinygrad.engine.realize import get_runner
 from tinygrad.nn.onnx import OnnxRunner
 
 OPENPILOT_MODEL = sys.argv[1] if len(sys.argv) > 1 else "https://github.com/commaai/openpilot/raw/v0.9.7/selfdrive/modeld/models/supercombo.onnx"
@@ -37,7 +36,7 @@ def compile(onnx_file):
     # copy i == 1 so use of JITBEAM is okay
     if i == 1: test_val = np.copy(ret)
   # iterate kernel CALLs in the captured LINEAR UOp; toposort descends into batched graph CUSTOM_FUNCTIONs
-  kernel_asts = {Ops.SINK, Ops.PROGRAM}
+  kernel_asts = {Ops.PROGRAM}
   kernel_calls = [u for u in run_onnx_jit.captured.linear.toposort(gate=lambda x: x.op not in kernel_asts)
                   if u.op is Ops.CALL and u.src[0].op in kernel_asts]
   print(f"captured {len(kernel_calls)} kernels")
@@ -49,8 +48,8 @@ def compile(onnx_file):
   read_image_count = 0
   gated_read_image_count = 0
   for call in kernel_calls:
-    device = next(b.device for b in call.src[1:] if b.op is not Ops.BIND)
-    src = get_runner(device, call.src[0]).p.src
+    _, _, _, source, _ = call.src[0].src
+    src = source.arg
     kernel_count += 1
     read_image_count += src.count("read_image")
     gated_read_image_count += src.count("?read_image")
@@ -134,14 +133,20 @@ def bench(run, inputs):
       run(**inputs).numpy()
 
 if __name__ == "__main__":
-  onnx_file = fetch(OPENPILOT_MODEL)
-  inputs, outputs = compile(onnx_file)
+  if getenv("RUN_PICKLE"):
+    with open(OUTPUT, "rb") as f: pickle_loaded = pickle.load(f)
+    inputs = {name: Tensor(Tensor.randn(*[int(s) for s in view.src[1].arg], dtype=dtype).numpy(), device=device)
+              for name, (view, _vars, dtype, device) in zip(pickle_loaded.captured.expected_names, pickle_loaded.captured.expected_input_info)}
+    test_vs_compile(pickle_loaded, inputs)
+  else:
+    onnx_file = fetch(OPENPILOT_MODEL)
+    inputs, outputs = compile(onnx_file)
 
-  with open(OUTPUT, "rb") as f: pickle_loaded = pickle.load(f)
+    with open(OUTPUT, "rb") as f: pickle_loaded = pickle.load(f)
 
-  test_vs_compile(pickle_loaded, inputs, outputs)
-  if getenv("SELFTEST"):
-    test_vs_onnx(inputs, outputs, onnx_file, 1e-4)
+    test_vs_compile(pickle_loaded, inputs, outputs)
+    if getenv("SELFTEST"):
+      test_vs_onnx(inputs, outputs, onnx_file, 1e-4)
 
   if getenv("BENCHMARK_LOG", ""):
     bench(pickle_loaded, inputs)

examples/webgpu/stable_diffusion/compile.py

@@ -114,7 +114,7 @@ if __name__ == "__main__":
     linear, output_bufs = jit_model(step, *step.input)
     functions, statements, bufs, _ = compile_net(linear, output_bufs)
     state = get_state_dict(model)
-    weights = {id(x.uop.base.realized): name for name, x in state.items()}
+    weights = {(id(b), b.offset, b.size, b.dtype): name for name, x in state.items() if (b:=x.uop.base.realized) is not None}
     kernel_code = '\n\n'.join([f"const {key} = `{fixup_code(code, key)}`;" for key, code in functions.items()])
     kernel_names = ', '.join([name for (name, _, _, _) in statements])
     input_names = [f"input{i}" for i in range(len(step.input))]

extra/amax/cast_amax.py

@@ -1,133 +0,0 @@
-import functools, pathlib
-from tinygrad import Tensor, dtypes
-from tinygrad.uop.ops import UOp, Ops, KernelInfo
-from tinygrad.renderer import Estimates
-from tinygrad.runtime.support.compiler_amd import HIPCCCompiler
-
-FP8_MAX = 448.0
-NUM_WG, THREADS_PER_WG = 1024, 256
-
-def _compile(cpp_name:str, n_elems:int, hidden:int):
-  src = (pathlib.Path(__file__).parent/cpp_name).read_text()
-  defines = [f"-DN_ELEMS={n_elems}", f"-DHIDDEN={hidden}", f"-DNUM_WG={NUM_WG}", f"-DTHREADS_PER_WG={THREADS_PER_WG}"]
-  return src, HIPCCCompiler("gfx950", ["-std=c++20", "-ffast-math", *defines]).compile_cached(src)
-
-def _shard_shape(shape:tuple, axis:int, ndev:int) -> list:
-  s = list(shape); s[axis] //= ndev; return s
-
-def _scalar_amax(amax_buf:Tensor) -> Tensor:
-  if isinstance(amax_buf.device, tuple):
-    from examples.mlperf.models.flat_llama import _local_abs_max
-    return _local_abs_max(amax_buf).detach()
-  return amax_buf.max().detach()
-
-# ** fused silu*mul -> fp8 cast + amax (w13 layout)
-
-@functools.cache
-def _custom_fused_bwd_w13(grad_xw13:UOp, xw13:UOp, grad_x2:UOp, amax_state:UOp, dname:str) -> UOp:
-  hidden = xw13.shape[2] // 2
-  n_elems = xw13.shape[0] * xw13.shape[1] * hidden
-  threads, workgroups = UOp.special(THREADS_PER_WG, "lidx0"), UOp.special(NUM_WG, "gidx0")
-  mem = n_elems * 2 * 5
-  sink = UOp.sink(grad_xw13.base, xw13.base, grad_x2.base, amax_state.base, threads, workgroups,
-                  arg=KernelInfo(f"fused_silu_mul_bwd_w13_{n_elems}", estimates=Estimates(ops=8*n_elems, mem=mem)))
-  src, lib = _compile("cast_amax_bwd_w13.cpp", n_elems, hidden)
-  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=(*sink.src, sink)),
-                               UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=lib)))
-
-@functools.cache
-def _custom_fused_cast_amax_w13(fp8_out:UOp, amax_buf:UOp, xw13:UOp, amax_state:UOp, dname:str) -> UOp:
-  hidden = xw13.shape[2] // 2
-  n_elems = xw13.shape[0] * xw13.shape[1] * hidden
-  threads, workgroups = UOp.special(THREADS_PER_WG, "lidx0"), UOp.special(NUM_WG, "gidx0")
-  mem = n_elems * 2 * 2 + n_elems + NUM_WG * 2
-  sink = UOp.sink(fp8_out.base, amax_buf.base, xw13.base, amax_state.base, threads, workgroups,
-                  arg=KernelInfo(f"fused_silu_mul_cast_amax_w13_{n_elems}", estimates=Estimates(ops=5*n_elems, mem=mem)))
-  src, lib = _compile("cast_amax_fwd_w13.cpp", n_elems, hidden)
-  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=(*sink.src, sink)),
-                               UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=lib)))
-
-def _fused_quantize_bwd_w13(gradient:UOp, kernel:UOp):
-  _, _, xw13, amax_state = kernel.src[1:]
-  device = xw13.device
-  if isinstance(device, tuple):
-    axis, ndev = xw13.axis, len(device)
-    assert axis in (0, 1), f"unsupported sharding axis={axis}"
-    grad_xw13 = Tensor(Tensor.invalids(*_shard_shape(xw13.shape, axis, ndev), dtype=dtypes.bfloat16, device=device).uop.multi(axis), device=device)
-    dname = device[0].split(":")[0]
-  else:
-    grad_xw13 = Tensor.invalids(*xw13.shape, dtype=dtypes.bfloat16, device=device)
-    dname = device.split(":")[0] if isinstance(device, str) else device
-  grad_x2_t = Tensor(gradient, device=device).cast(dtypes.bfloat16)
-  fxn = functools.partial(_custom_fused_bwd_w13, dname=dname)
-  grad_xw13, *_ = Tensor.custom_kernel(grad_xw13, Tensor(xw13, device=device), grad_x2_t, Tensor(amax_state, device=device), fxn=fxn)
-  return (None, None, grad_xw13.uop, None)
-
-def fused_quantize_fp8_w13(xw13:Tensor, amax_state:Tensor, fp8_dtype) -> tuple[Tensor, Tensor, Tensor]:
-  # silu(xw1)*xw3 -> fp8 + amax over fused xw13 layout. Returns (fp8, inv_scale, new_amax).
-  assert xw13.dtype == dtypes.bfloat16, f"expected bf16, got {xw13.dtype}"
-  MBS, SEQ, H2 = xw13.shape
-  assert H2 % 2 == 0, f"w13 last-axis must be even, got {H2}"
-  HIDDEN = H2 // 2
-  if isinstance(xw13.device, tuple):
-    axis, ndev = xw13.uop.axis, len(xw13.device)
-    assert axis in (0, 1), f"unsupported sharding axis={axis}"
-    fp8_out = Tensor(Tensor.invalids(*_shard_shape((MBS, SEQ, HIDDEN), axis, ndev), dtype=fp8_dtype, device=xw13.device).uop.multi(axis), device=xw13.device)
-    amax_buf = Tensor(Tensor.invalids(NUM_WG, dtype=dtypes.bfloat16, device=xw13.device).uop.multi(0), device=xw13.device)
-    dname = xw13.device[0].split(":")[0]
-  else:
-    fp8_out = Tensor.invalids(MBS, SEQ, HIDDEN, dtype=fp8_dtype, device=xw13.device)
-    amax_buf = Tensor.invalids(NUM_WG, dtype=dtypes.bfloat16, device=xw13.device)
-    dname = xw13.device.split(":")[0] if isinstance(xw13.device, str) else xw13.device
-  fxn = functools.partial(_custom_fused_cast_amax_w13, dname=dname)
-  fp8_out, amax_buf, *_ = Tensor.custom_kernel(fp8_out, amax_buf, xw13, amax_state, fxn=fxn, grad_fxn=_fused_quantize_bwd_w13)
-  inv_scale = (amax_state.float() + 1e-8) / FP8_MAX
-  return fp8_out, inv_scale, _scalar_amax(amax_buf)
-
-# ** fused (x * weight) -> fp8 cast + amax (norm-mul-quantize)
-
-@functools.cache
-def _custom_mul_quantize_fp8(fp8_out:UOp, amax_buf:UOp, x:UOp, weight:UOp, amax_state:UOp, dname:str) -> UOp:
-  MBS, SEQ, HIDDEN = x.shape
-  n_elems = MBS * SEQ * HIDDEN
-  threads, workgroups = UOp.special(THREADS_PER_WG, "lidx0"), UOp.special(NUM_WG, "gidx0")
-  mem = n_elems * 2 + HIDDEN * 2 + n_elems + NUM_WG * 2
-  sink = UOp.sink(fp8_out.base, amax_buf.base, x.base, weight.base, amax_state.base, threads, workgroups,
-                  arg=KernelInfo(f"fused_mul_quantize_fp8_{n_elems}_h{HIDDEN}", estimates=Estimates(ops=3*n_elems, mem=mem)))
-  src, lib = _compile("fused_mul_quantize_fp8.cpp", n_elems, HIDDEN)
-  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=(*sink.src, sink)),
-                               UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=lib)))
-
-def _fused_mul_quantize_fp8_bwd(gradient:UOp, kernel:UOp):
-  # inputs: (fp8_out, amax_buf, x, weight, amax_state); grads for x and weight
-  _, _, x_u, weight_u, amax_state_u = kernel.src[1:]
-  device = x_u.device
-  grad_t = Tensor(gradient, device=device).cast(dtypes.bfloat16)
-  x_t, weight_t = Tensor(x_u, device=device), Tensor(weight_u, device=device)
-  scale = FP8_MAX / (Tensor(amax_state_u, device=device).float() + 1e-8)
-  grad_scaled = grad_t.float() * scale
-  # grad_x stays bf16 to avoid CSE materializing a (MBS, SEQ, HIDDEN) fp32 intermediate
-  grad_x = (grad_scaled * weight_t.float()).cast(dtypes.bfloat16)
-  grad_weight = (grad_scaled * x_t.float()).sum(axis=(0, 1)).cast(dtypes.bfloat16)
-  return (None, None, grad_x.uop, grad_weight.uop, None)
-
-def fused_mul_quantize_fp8(x:Tensor, weight:Tensor, amax_state:Tensor, fp8_dtype) -> tuple[Tensor, Tensor, Tensor]:
-  # (x * weight) -> fp8 + amax, delayed scaling. Returns (fp8, inv_scale, new_amax).
-  assert x.dtype == dtypes.bfloat16 and weight.dtype == dtypes.bfloat16
-  assert x.shape[-1] == weight.shape[-1], f"HIDDEN mismatch: x={x.shape}, weight={weight.shape}"
-  MBS, SEQ, HIDDEN = x.shape
-  if isinstance(x.device, tuple):
-    axis, ndev = x.uop.axis, len(x.device)
-    assert axis in (0, 1), f"unsupported sharding axis={axis}"
-    fp8_out = Tensor(Tensor.invalids(*_shard_shape((MBS, SEQ, HIDDEN), axis, ndev), dtype=fp8_dtype, device=x.device).uop.multi(axis), device=x.device)
-    amax_buf = Tensor(Tensor.invalids(NUM_WG, dtype=dtypes.bfloat16, device=x.device).uop.multi(0), device=x.device)
-    dname = x.device[0].split(":")[0]
-  else:
-    fp8_out = Tensor.invalids(MBS, SEQ, HIDDEN, dtype=fp8_dtype, device=x.device)
-    amax_buf = Tensor.invalids(NUM_WG, dtype=dtypes.bfloat16, device=x.device)
-    dname = x.device.split(":")[0] if isinstance(x.device, str) else x.device
-  fxn = functools.partial(_custom_mul_quantize_fp8, dname=dname)
-  fp8_out, amax_buf, *_ = Tensor.custom_kernel(fp8_out, amax_buf, x, weight, amax_state, fxn=fxn, grad_fxn=_fused_mul_quantize_fp8_bwd)
-  new_amax = _scalar_amax(amax_buf)
-  inv_scale = (amax_state.float() + 1e-8) / FP8_MAX
-  return fp8_out, inv_scale, new_amax

extra/export_model.py

@@ -6,7 +6,7 @@ from tinygrad.engine.jit import TinyJit
 from tinygrad.nn.state import get_state_dict
 from tinygrad.helpers import Context, to_mv, prod
 from tinygrad.uop.ops import Ops, UOp
-from tinygrad.codegen import get_program
+from tinygrad.codegen import to_program
 import json
 from collections import OrderedDict
 
@@ -36,10 +36,11 @@ def compile_net(linear:UOp, output_bufs:List[Buffer]) -> Tuple[Dict[str,str], Li
 
   for call in iter_kernel_calls(linear):
     arg_uops = [b for b in call.src[1:] if b.op is not Ops.BIND]
-    prg = get_program(call.src[0], Device[arg_uops[0].device].renderer)
-    functions[prg.function_name] = prg.src
-    cargs = [name_of(bu, i == 0) for i, bu in enumerate(arg_uops)] + [v for v in prg.vars if v.op is Ops.DEFINE_VAR]
-    statements.append((prg.function_name, cargs, prg.global_size, prg.local_size))
+    prg = to_program(call.src[0], Device[arg_uops[0].device].renderer)
+    info = prg.arg
+    functions[info.function_name] = prg.src[3].arg
+    cargs = [name_of(bu, i == 0) for i, bu in enumerate(arg_uops)] + [v for v in info.vars if v.op is Ops.DEFINE_VAR]
+    statements.append((info.function_name, cargs, info.global_size, info.local_size))
 
   return functions, statements, {name:(size, dtype, key) for name, size, dtype, key in bufs.values()}, bufs_to_save
 
@@ -244,7 +245,7 @@ def export_model(model, target:str, *inputs, model_name: Optional[str] = "model"
   with Context(JIT=2, CPU_COUNT=1): linear, output_bufs = jit_model(model, *inputs)
   functions, statements, bufs, bufs_to_save = compile_net(linear, output_bufs)
   state = get_state_dict(model)
-  weight_names = {id(x.uop.base.realized): name for name, x in state.items()}
+  weight_names = {(id(b), b.offset, b.size, b.dtype): name for name, x in state.items() if (b:=x.uop.base.realized) is not None}
   input_names = [f"input{i}" for i in range(len(inputs))]
   output_names = [f"output{i}" for i in range(len(output_bufs))]
 

extra/gemm/amd_asm_matmul.py

@@ -13,7 +13,7 @@ from tinygrad import Tensor, Device, Context, GlobalCounters
 from tinygrad.uop.ops import UOp, Ops, KernelInfo
 from tinygrad.helpers import getenv, colored
 from tinygrad.dtype import dtypes, AddrSpace
-from tinygrad.engine.realize import Estimates
+from tinygrad.engine.realize import Estimates, run_linear
 from tinygrad.renderer.amd.dsl import s, v, VCC_LO, NULL
 from tinygrad.runtime.autogen.amd.rdna3.ins import *
 
@@ -463,11 +463,14 @@ def test_matmul():
                                                                                   estimates=Estimates(ops=N*N*N*2, mem=N*N*4*3)))
     return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))
   c = Tensor.custom_kernel(a, b, c, fxn=asm_kernel)[2]
-  ei = c.schedule()[0].lower()
+  linear = c.schedule_linear()
 
   ets = []
   with Context(DEBUG=2):
-    for _ in range(getenv("CNT", 5)): ets.append(ei.run(wait=True))
+    for _ in range(getenv("CNT", 5)):
+      start = GlobalCounters.time_sum_s
+      run_linear(linear)
+      ets.append(GlobalCounters.time_sum_s - start)
   print(f"REAL TFLOPS {N * N * N * 2 / min(ets) * 1e-12:.2f}")
 
   if getenv("VERIFY", 1):

extra/gemm/amd_matmul.py

@@ -1,31 +1,39 @@
 # kernel8_batched_gmem.s from https://seb-v.github.io/optimization/update/2025/01/20/Fast-GPU-Matrix-multiplication.html
 # sudo PATH=/opt/homebrew/Cellar/llvm/20.1.6/bin:$PATH AMD_LLVM=0 AMD=1 DEBUG=2 python3 extra/gemm/amd_matmul.py
 import pathlib
-from dataclasses import replace
 from tinygrad import Tensor, Device, Context, GlobalCounters
 from tinygrad.helpers import getenv
-from tinygrad.engine.realize import CompiledRunner, ExecItem, get_program
+from tinygrad.uop.ops import UOp, Ops, KernelInfo
+from tinygrad.renderer import Estimates
+from tinygrad.engine.realize import run_linear
 
 N = 4096
 run_count = 5
 
-if __name__ == "__main__":
-  ast = (Tensor.empty(N, N)@Tensor.empty(N, N)).schedule()[-1].ast
-  prg = get_program(ast, Device.default.renderer)
+def make_matmul_kernel(name:str, src:str, local_size:int):
+  def fxn(a:UOp, b:UOp, c:UOp) -> UOp:
+    threads = UOp.special(local_size, "lidx0")
+    wg_x = UOp.special(N//128, "gidx0")
+    wg_y = UOp.special(N//128, "gidx1")
+    sink = UOp.sink(a.base, b.base, c.base, threads, wg_x, wg_y, arg=KernelInfo(name, estimates=Estimates(ops=2*N**3, mem=3*N*N*4)))
+    lib = Device[Device.DEFAULT].compiler.compile_cached(src)
+    return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=Device.DEFAULT), UOp(Ops.LINEAR, src=(*sink.src, sink)),
+                                 UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=lib)))
+  return fxn
 
+if __name__ == "__main__":
   if getenv("ASM") == 1:
     src = (pathlib.Path(__file__).parent / "amd_seb" / "kernel8_batched_gmem.s").read_text()
-    prgfast = replace(prg, name="kernel", src=src, global_size=[N//128, N//128, 1], local_size=[128, 1, 1])
+    name, local_size = "kernel", 128
   elif getenv("ASM") == -1:
     src = (pathlib.Path(__file__).parent / "amd_seb" / "kernel3_registers.cpp").read_text()
-    prgfast = replace(prg, name="kernel3_registers", src=src, global_size=[N//128, N//128, 1], local_size=[256, 1, 1])
+    name, local_size = "kernel3_registers", 256
   elif getenv("ASM") == -2:
     src = (pathlib.Path(__file__).parent / "amd_seb" / "kernel4_gmem_df.cpp").read_text()
-    prgfast = replace(prg, name="kernel4_gmem_db", src=src, global_size=[N//128, N//128, 1], local_size=[256, 1, 1])
+    name, local_size = "kernel4_gmem_db", 256
   else:
     src = (pathlib.Path(__file__).parent / "amd_seb" / "kernel5_lds_optim.cpp").read_text()
-    prgfast = replace(prg, name="kernel5_lds_optim", src=src, global_size=[N//128, N//128, 1], local_size=[128, 1, 1])
-  runner = CompiledRunner(prgfast)
+    name, local_size = "kernel5_lds_optim", 128
 
   a = Tensor.randn(N, N).realize()
   b = Tensor.randn(N, N).realize()
@@ -35,8 +43,8 @@ if __name__ == "__main__":
   with Context(DEBUG=2):
     for _ in range(run_count): tc = (a@b).realize()
 
+  linear = Tensor.custom_kernel(a, b, c, fxn=make_matmul_kernel(name, src, local_size))[2].schedule_linear()
   GlobalCounters.reset()
-  ei = ExecItem(ast, [a.uop.buffer, b.uop.buffer, c.uop.buffer], prg=runner)
   with Context(DEBUG=2):
-    for _ in range(run_count): ei.run(wait=True)
+    for _ in range(run_count): run_linear(linear)
   print(f"custom  {(c-tc).square().mean().item()}")

extra/gemm/max_matmul.py

@@ -1,7 +1,6 @@
 import numpy as np, os
 from tinygrad.helpers import getenv, flat_mv
 from tinygrad import dtypes
-from tinygrad.engine.realize import get_program
 
 # for copied uops
 from tinygrad import dtypes

extra/gemm/rdna4_asm_matmul.py

@@ -4,7 +4,7 @@ from tinygrad import Tensor, Device, Context, GlobalCounters
 from tinygrad.uop.ops import UOp, Ops, KernelInfo
 from tinygrad.helpers import getenv, colored
 from tinygrad.dtype import dtypes, AddrSpace
-from tinygrad.engine.realize import Estimates
+from tinygrad.engine.realize import Estimates, run_linear
 from tinygrad.renderer.amd.dsl import s, v, VCC_LO, NULL, src, ttmp
 from tinygrad.runtime.autogen.amd.rdna4.ins import *
 
@@ -225,11 +225,14 @@ def test_matmul():
     return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))
 
   c = Tensor.custom_kernel(a, b, c, fxn=asm_kernel)[2]
-  ei = c.schedule()[0].lower()
+  linear = c.schedule_linear()
 
   ets = []
   with Context(DEBUG=2):
-    for _ in range(getenv("CNT", 5)): ets.append(ei.run(wait=True))
+    for _ in range(getenv("CNT", 5)):
+      start = GlobalCounters.time_sum_s
+      run_linear(linear)
+      ets.append(GlobalCounters.time_sum_s - start)
   print(f"REAL TFLOPS {N*N*N*2 / min(ets) * 1e-12:.2f}")
 
   if getenv("VERIFY", 1):

extra/gemm/simple_matmul.py

@@ -2,6 +2,7 @@ import numpy as np
 from tinygrad import dtypes, Tensor
 from tinygrad.helpers import getenv, get_single_element
 from tinygrad.dtype import _to_np_dtype
+from tinygrad.engine.realize import compile_linear
 from tinygrad.codegen.opt import OptOps
 
 dtype_in = (dtypes.half if getenv("HALF") else dtypes.bfloat16 if getenv("BFLOAT16") else
@@ -38,10 +39,10 @@ if __name__ == "__main__":
     c = a.matmul(b, dtype=acc_dtype).realize()
 
   if getenv("SHOULD_USE_TC"):
-    sched = a.matmul(b, dtype=acc_dtype).schedule()
-    ei = get_single_element(sched)
-    ei.lower()
-    assert any(opt.op is OptOps.TC for opt in ei.prg.p.applied_opts), f"TC not triggered, {ei.prg.p.applied_opts}"
+    linear = compile_linear(a.matmul(b, dtype=acc_dtype).schedule_linear())
+    call = get_single_element(list(linear.src))
+    applied_opts = call.src[0].src[0].arg.applied_opts
+    assert any(opt.op is OptOps.TC for opt in applied_opts), f"TC not triggered, {applied_opts}"
 
   ref = a.numpy().astype(np.float32) @ b.numpy().astype(np.float32)
   res = c.numpy()

extra/gemm/tinygrad_nv_matmul.py

@@ -1,7 +1,7 @@
-from tinygrad import Tensor, dtypes, Device
-from tinygrad.helpers import getenv, DEBUG
-from tinygrad.codegen.opt.kernel import Kernel, Opt, OptOps
-from tinygrad.engine.realize import CompiledRunner, ExecItem, get_program
+from tinygrad import Tensor, dtypes, Context
+from tinygrad.helpers import getenv
+from tinygrad.codegen.opt import Opt, OptOps
+from tinygrad.engine.realize import run_linear
 from dataclasses import replace
 
 N = 4096
@@ -11,9 +11,6 @@ if __name__ == "__main__":
   else:
     A, B = Tensor.empty(N, N, dtype=dtypes.float16), Tensor.empty(N, N, dtype=dtypes.float16)
   C = A.matmul(B)
-  si = C.schedule()[-1]
-  ast = si.ast
-  k = Kernel(ast, opts=Device[Device.DEFAULT].renderer)
   if getenv("GEMV"):
     opts = [
       Opt(op=OptOps.UNROLL, axis=0, amt=8),
@@ -28,10 +25,10 @@ if __name__ == "__main__":
       Opt(op=OptOps.LOCAL, axis=1, amt=2),
       Opt(op=OptOps.LOCAL, axis=0, amt=2),
     ]
-  k.apply_opts(opts)
-  prg = get_program(k.ast.replace(arg=replace(k.ast.arg, opts_to_apply=tuple(k.applied_opts))), k.opts)
-  new_src = prg.src
-  # can mod source here
-  prg = replace(prg, src=new_src)
-  ei = ExecItem(si.ast, [x.ensure_allocated() for x in si.bufs], si.metadata, prg=CompiledRunner(prg))
-  for i in range(5): ei.run(wait=True)
+  linear = C.schedule_linear()
+  call = linear.src[-1]
+  new_ast = call.src[0].replace(arg=replace(call.src[0].arg, opts_to_apply=tuple(opts)))
+  new_call = call.replace(src=(new_ast, *call.src[1:]))
+  linear = linear.replace(src=tuple(new_call if c is call else c for c in linear.src))
+  with Context(DEBUG=2):
+    for i in range(5): run_linear(linear)

extra/gemm/triton_nv_matmul.py

@@ -4,7 +4,8 @@ import triton.language as tl
 from triton.compiler import AttrsDescriptor, ASTSource, compile as triton_compile
 import numpy as np
 from tinygrad import Tensor, dtypes, Device
-from tinygrad.engine.realize import CompiledRunner, ExecItem, ProgramSpec
+from tinygrad.engine.realize import CompiledRunner
+from tinygrad.uop.ops import Ops, UOp, KernelInfo, ProgramInfo
 from tinygrad.helpers import getenv
 np.set_printoptions(suppress=True)
 
@@ -73,9 +74,11 @@ if __name__ == "__main__":
 
   A, B = Tensor.normal(M, K, std=1e-1, dtype=dtypes.float16).realize(), Tensor.normal(K, N, std=1e-1, dtype=dtypes.float16).realize()
   C = A.matmul(B)
-  from tinygrad.schedule import linear_to_schedule
+  from tinygrad.uop.ops import Ops
   linear, var_vals = C.linear_with_vars()
-  si = linear_to_schedule(linear)[-1]
+  last_call = linear.src[-1]
+  ast = last_call.src[0]
+  bufs = [s.buffer for s in last_call.src[1:] if s.op is not Ops.BIND]
 
   src = compiled.asm["ptx"]
   # specify the shared memory here so we don't need to do it dynamically
@@ -86,13 +89,16 @@ if __name__ == "__main__":
   # remove debug sections
   src = src.split("\t.file")[0]
   assert '.extern .shared' not in src
-  prg = ProgramSpec("matmul_kernel", src, device=Device.DEFAULT,
-                global_size=[M//BLOCK_SIZE_M, N//BLOCK_SIZE_N, 1], local_size=[32*compiled.metadata.num_warps, 1, 1],
-                mem_estimate=A.nbytes() + B.nbytes() + C.nbytes())
-  ei = ExecItem(si.ast, [x.ensure_allocated() for x in si.bufs], si.metadata, prg=CompiledRunner(prg))
+  info = ProgramInfo(name="matmul_kernel",
+                     global_size=(M//BLOCK_SIZE_M, N//BLOCK_SIZE_N, 1), local_size=(32*compiled.metadata.num_warps, 1, 1))
+  sink = UOp.sink(arg=KernelInfo(name="matmul_kernel"))
+  prg_uop = UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=Device.DEFAULT), UOp(Ops.LINEAR), UOp(Ops.SOURCE, arg=src)), arg=info)
+  runner = CompiledRunner(prg_uop, Device.DEFAULT)
+  all_bufs = [x.ensure_allocated() for x in bufs]
+  prg_bufs = [all_bufs[i] for i in runner.p.globals]
   tflops = []
   for i in range(5):
-    tm = ei.run(wait=True)
+    tm = runner(prg_bufs, {}, wait=True)
     tflops.append((2*M*K*N/tm)*1e-12)
   print(f"TFLOPS: {max(tflops):.2f}")
 

extra/gemm/tvm_gemm.py

@@ -36,10 +36,10 @@ A = Tensor.rand(M, K, device="CPU")
 B = Tensor.rand(K, N, device="CPU")
 C = (A.reshape(M, 1, K) * B.permute(1,0).reshape(1, N, K)).sum(axis=2)
 
-sched = C.schedule()
+linear = C.schedule_linear()
 from tinygrad.codegen.opt.kernel import Kernel
 from tinygrad.device import CompilerOptions
-lin = Kernel(sched[-1].ast, CompilerOptions(has_local=False, supports_float4=False))
+lin = Kernel(linear.src[-1].src[0], CompilerOptions(has_local=False, supports_float4=False))
 lin.to_program()
 from tinygrad.runtime.ops_cpu import renderer
 src = renderer("mmult", lin.uops)

extra/llama_kernels/__init__.py

@@ -0,0 +1,35 @@
+from __future__ import annotations
+import functools, pathlib
+from tinygrad import Tensor, dtypes
+from tinygrad.uop.ops import Ops
+from tinygrad.runtime.support.compiler_amd import HIPCCCompiler
+
+FP8_MAX = 448.0
+NUM_WG, THREADS_PER_WG = 1024, 256
+
+# per-device abs max without allreduce
+@functools.cache
+def _local_abs_max_fxn(x_p, device):
+  x = Tensor(x_p, device=device)
+  inner = Tensor(x.uop.src[0]) if x.uop.op is Ops.MULTI else x
+  return (inner.abs().max(),)
+
+def local_abs_max(x:Tensor) -> Tensor:
+  param = x.as_param(0)
+  fxn = _local_abs_max_fxn(param.uop, x.device)
+  return Tensor(fxn[0].uop.call(x.uop).gettuple(0))
+
+def scalar_amax(amax_buf:Tensor) -> Tensor:
+  if isinstance(amax_buf.device, tuple):
+    return local_abs_max(amax_buf).detach()
+  return amax_buf.max().detach()
+
+def shard_shape(shape:tuple, axis:int, ndev:int) -> list:
+  s = list(shape)
+  s[axis] //= ndev
+  return s
+
+def compile_cpp(cpp_dir:pathlib.Path, cpp_name:str, n_elems:int, hidden:int):
+  src = (cpp_dir/cpp_name).read_text()
+  defines = [f"-DN_ELEMS={n_elems}", f"-DHIDDEN={hidden}", f"-DNUM_WG={NUM_WG}", f"-DTHREADS_PER_WG={THREADS_PER_WG}"]
+  return src, HIPCCCompiler("gfx950", ["-std=c++20", "-ffast-math", *defines]).compile_cached(src)

extra/llama_kernels/cast_amax/__init__.py

@@ -0,0 +1,73 @@
+from __future__ import annotations
+import functools, pathlib
+from tinygrad import Tensor, dtypes
+from tinygrad.uop.ops import UOp, Ops, KernelInfo
+from tinygrad.renderer import Estimates
+from extra.llama_kernels import FP8_MAX, NUM_WG, THREADS_PER_WG, compile_cpp, shard_shape, scalar_amax
+
+@functools.cache
+def _custom_fused_bwd_w13(grad_xw13:UOp, xw13:UOp, grad_x2:UOp, amax_state:UOp, dname:str) -> UOp:
+  hidden = xw13.shape[2] // 2
+  n_elems = xw13.shape[0] * xw13.shape[1] * hidden
+  threads, workgroups = UOp.special(THREADS_PER_WG, "lidx0"), UOp.special(NUM_WG, "gidx0")
+  mem = n_elems * 2 * 5
+  sink = UOp.sink(grad_xw13.base, xw13.base, grad_x2.base, amax_state.base, threads, workgroups,
+                  arg=KernelInfo(f"fused_silu_mul_bwd_w13_{n_elems}", estimates=Estimates(ops=8*n_elems, mem=mem)))
+  src, lib = compile_cpp(pathlib.Path(__file__).parent, "cast_amax_bwd_w13.cpp", n_elems, hidden)
+  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=(*sink.src, sink)),
+                               UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=lib)))
+
+@functools.cache
+def _custom_fused_cast_amax_w13(fp8_out:UOp, amax_buf:UOp, xw13:UOp, amax_state:UOp, dname:str) -> UOp:
+  hidden = xw13.shape[2] // 2
+  n_elems = xw13.shape[0] * xw13.shape[1] * hidden
+  threads, workgroups = UOp.special(THREADS_PER_WG, "lidx0"), UOp.special(NUM_WG, "gidx0")
+  mem = n_elems * 2 * 2 + n_elems + NUM_WG * 2
+  sink = UOp.sink(fp8_out.base, amax_buf.base, xw13.base, amax_state.base, threads, workgroups,
+                  arg=KernelInfo(f"fused_silu_mul_cast_amax_w13_{n_elems}", estimates=Estimates(ops=5*n_elems, mem=mem)))
+  src, lib = compile_cpp(pathlib.Path(__file__).parent, "cast_amax_fwd_w13.cpp", n_elems, hidden)
+  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=(*sink.src, sink)),
+                               UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=lib)))
+
+def _fused_quantize_bwd_w13(gradient:UOp, kernel:UOp):
+  # NOTE: inputs are (fp8_out, amax_buf, xw13, amax_state); grad for xw13 only
+  _, _, xw13, amax_state = kernel.src[1:]
+  device = xw13.device
+  if isinstance(device, tuple):
+    axis, ndev = xw13.axis, len(device)
+    assert axis in (0, 1), f"unsupported sharding axis={axis}"
+    grad_xw13 = Tensor(Tensor.invalids(*shard_shape(xw13.shape, axis, ndev), dtype=dtypes.bfloat16,
+                                        device=device).uop.multi(axis), device=device)
+    dname = device[0].split(":")[0]
+  else:
+    grad_xw13 = Tensor.invalids(*xw13.shape, dtype=dtypes.bfloat16, device=device)
+    dname = device.split(":")[0] if isinstance(device, str) else device
+  grad_x2_t = Tensor(gradient, device=device).cast(dtypes.bfloat16)
+  fxn = functools.partial(_custom_fused_bwd_w13, dname=dname)
+  grad_xw13, *_ = Tensor.custom_kernel(grad_xw13, Tensor(xw13, device=device), grad_x2_t,
+                                        Tensor(amax_state, device=device), fxn=fxn)
+  return (None, None, grad_xw13.uop, None)
+
+def fused_quantize_fp8_w13(xw13:Tensor, amax_state:Tensor, fp8_dtype) -> tuple[Tensor, Tensor, Tensor]:
+  # NOTE: silu(xw1)*xw3 -> fp8 + amax over fused xw13 layout. Returns (fp8, inv_scale, new_amax)
+  assert xw13.dtype == dtypes.bfloat16, f"expected bf16, got {xw13.dtype}"
+  MBS, SEQ, H2 = xw13.shape
+  assert H2 % 2 == 0, f"w13 last-axis must be even, got {H2}"
+  HIDDEN = H2 // 2
+  if isinstance(xw13.device, tuple):
+    axis, ndev = xw13.uop.axis, len(xw13.device)
+    assert axis in (0, 1), f"unsupported sharding axis={axis}"
+    fp8_out = Tensor(Tensor.invalids(*shard_shape((MBS, SEQ, HIDDEN), axis, ndev), dtype=fp8_dtype,
+                                     device=xw13.device).uop.multi(axis), device=xw13.device)
+    amax_buf = Tensor(Tensor.invalids(NUM_WG, dtype=dtypes.bfloat16, device=xw13.device).uop.multi(0),
+                      device=xw13.device)
+    dname = xw13.device[0].split(":")[0]
+  else:
+    fp8_out = Tensor.invalids(MBS, SEQ, HIDDEN, dtype=fp8_dtype, device=xw13.device)
+    amax_buf = Tensor.invalids(NUM_WG, dtype=dtypes.bfloat16, device=xw13.device)
+    dname = xw13.device.split(":")[0] if isinstance(xw13.device, str) else xw13.device
+  fxn = functools.partial(_custom_fused_cast_amax_w13, dname=dname)
+  fp8_out, amax_buf, *_ = Tensor.custom_kernel(fp8_out, amax_buf, xw13, amax_state, fxn=fxn,
+                                                grad_fxn=_fused_quantize_bwd_w13)
+  inv_scale = (amax_state.float() + 1e-8) / FP8_MAX
+  return fp8_out, inv_scale, scalar_amax(amax_buf)

extra/llama_kernels/fused_ce/__init__.py

@@ -0,0 +1,98 @@
+from __future__ import annotations
+import functools, pathlib
+from tinygrad import Tensor, dtypes
+from tinygrad.uop.ops import UOp, Ops, KernelInfo
+from tinygrad.renderer import Estimates
+from tinygrad.runtime.support.compiler_amd import HIPCCCompiler
+
+THREADS_PER_WG = 256
+
+@functools.cache
+def _custom_fused_ce_loss_fwd(loss_out:UOp, max_out:UOp, lse_out:UOp, logits:UOp, targets:UOp,
+                              dname:str, vocab:int, rows:int, label_smoothing:float) -> UOp:
+  threads, workgroups = UOp.special(THREADS_PER_WG, "lidx0"), UOp.special(rows, "gidx0")
+  mem = rows * vocab * 2 + rows * 12 + rows * 4
+  sink = UOp.sink(loss_out.base, max_out.base, lse_out.base, logits.base, targets.base,
+                  threads, workgroups,
+                  arg=KernelInfo(f"fused_ce_loss_fwd", estimates=Estimates(ops=6*rows*vocab, mem=mem)))
+  src = (pathlib.Path(__file__).parent/"fused_ce_loss.cpp").read_text()
+  defines = [f"-DVOCAB={vocab}", f"-DTHREADS_PER_WG={THREADS_PER_WG}",
+             f"-DLABEL_SMOOTHING={label_smoothing}f"]
+  lib = HIPCCCompiler("gfx950", ["-std=c++20", "-ffast-math", *defines]).compile_cached(src)
+  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=(*sink.src, sink)),
+                               UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=lib)))
+
+@functools.cache
+def _custom_fused_ce_loss_bwd(d_logits:UOp, logits:UOp, lse:UOp, targets:UOp, scale:UOp,
+                              dname:str, vocab:int, rows:int, label_smoothing:float) -> UOp:
+  threads, workgroups = UOp.special(THREADS_PER_WG, "lidx0"), UOp.special(rows, "gidx0")
+  mem = rows * vocab * 4 + rows * 8 + 4
+  sink = UOp.sink(d_logits.base, logits.base, lse.base, targets.base, scale.base,
+                  threads, workgroups,
+                  arg=KernelInfo(f"fused_ce_loss_bwd", estimates=Estimates(ops=4*rows*vocab, mem=mem)))
+  src = (pathlib.Path(__file__).parent/"fused_ce_loss_bwd.cpp").read_text()
+  defines = [f"-DVOCAB={vocab}", f"-DTHREADS_PER_WG={THREADS_PER_WG}",
+             f"-DLABEL_SMOOTHING={label_smoothing}f"]
+  lib = HIPCCCompiler("gfx950", ["-std=c++20", "-ffast-math", *defines]).compile_cached(src)
+  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=(*sink.src, sink)),
+                               UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=lib)))
+
+def _fused_ce_loss_bwd(gradient:UOp, kernel:UOp, label_smoothing:float):
+  # NOTE: forward inputs are (loss_out, max_out, lse_out, logits, targets)
+  # gradient is the upstream grad w.r.t. per-row loss (shape: (rows,) fp32)
+  _, _, lse_u, logits_u, targets_u = kernel.src[1:]
+  device = logits_u.device
+  rows_vocab = logits_u.shape  # (rows, VOCAB) after reshape
+  rows, VOCAB = rows_vocab
+  if isinstance(device, tuple):
+    axis = logits_u.axis
+    ndev = len(device)
+    d_logits = Tensor(Tensor.invalids(rows // ndev, VOCAB, dtype=dtypes.bfloat16, device=device).uop.multi(axis), device=device)
+    dname = device[0].split(":")[0]
+    rows_per_dev = rows // ndev
+  else:
+    d_logits = Tensor.invalids(rows, VOCAB, dtype=dtypes.bfloat16, device=device)
+    dname = device.split(":")[0] if isinstance(device, str) else device
+    rows_per_dev = rows
+  grad_t = Tensor(gradient, device=device).float().reshape(-1)  # (rows,) fp32
+  # NOTE: .mean() backward gives same grad per row (1/N), so broadcast is safe; take scalar
+  scale = grad_t[0:1].contiguous()
+  logits_t = Tensor(logits_u.after(kernel), device=device)
+  lse_t = Tensor(lse_u.after(kernel), device=device)
+  targets_t = Tensor(targets_u, device=device)
+  fxn = functools.partial(_custom_fused_ce_loss_bwd, dname=dname, vocab=VOCAB, rows=rows_per_dev, label_smoothing=label_smoothing)
+  d_logits, *_ = Tensor.custom_kernel(d_logits, logits_t, lse_t, targets_t, scale, fxn=fxn)
+  return (None, None, None, d_logits.uop, None)
+
+def fused_ce_loss(logits:Tensor, targets:Tensor, label_smoothing:float=0.1) -> Tensor:
+  # NOTE: fused sparse_categorical_crossentropy with label smoothing, returns mean loss scalar
+  assert logits.dtype == dtypes.bfloat16, f"expected bf16, got {logits.dtype}"
+  assert logits.ndim == 3, f"expected (MBS, SEQ, VOCAB), got {logits.shape}"
+  MBS, SEQ, VOCAB = logits.shape
+  rows = MBS * SEQ
+  if isinstance(logits.device, tuple):
+    axis = logits.uop.axis
+    assert axis in (0, 1), f"unsupported sharding axis={axis} for CE loss"
+    ndev = len(logits.device)
+    loss_out = Tensor(Tensor.invalids(rows // ndev, dtype=dtypes.float32, device=logits.device).uop.multi(0),
+                      device=logits.device)
+    max_out  = Tensor(Tensor.invalids(rows // ndev, dtype=dtypes.float32, device=logits.device).uop.multi(0),
+                      device=logits.device)
+    lse_out  = Tensor(Tensor.invalids(rows // ndev, dtype=dtypes.float32, device=logits.device).uop.multi(0),
+                      device=logits.device)
+    dname = logits.device[0].split(":")[0]
+    rows_per_dev = rows // ndev
+  else:
+    loss_out = Tensor.invalids(rows, dtype=dtypes.float32, device=logits.device)
+    max_out  = Tensor.invalids(rows, dtype=dtypes.float32, device=logits.device)
+    lse_out  = Tensor.invalids(rows, dtype=dtypes.float32, device=logits.device)
+    dname = logits.device.split(":")[0] if isinstance(logits.device, str) else logits.device
+    rows_per_dev = rows
+  logits_flat = logits.reshape(rows, VOCAB)
+  targets_flat = targets.reshape(-1).cast(dtypes.int32)
+  fxn = functools.partial(_custom_fused_ce_loss_fwd, dname=dname, vocab=VOCAB, rows=rows_per_dev,
+                          label_smoothing=label_smoothing)
+  loss_out, max_out, lse_out, *_ = Tensor.custom_kernel(
+    loss_out, max_out, lse_out, logits_flat, targets_flat,
+    fxn=fxn, grad_fxn=functools.partial(_fused_ce_loss_bwd, label_smoothing=label_smoothing))
+  return loss_out.mean()

extra/llama_kernels/fused_ce/fused_ce_loss.cpp

@@ -0,0 +1,104 @@
+#include <hip/hip_runtime.h>
+#include <hip/hip_bf16.h>
+
+// Fused forward sparse-CE with label smoothing.
+// SINGLE-PASS online softmax + vectorized 8-wide bf16 loads for HBM coalescing.
+
+#ifndef VOCAB
+#define VOCAB 128256
+#endif
+#ifndef THREADS_PER_WG
+#define THREADS_PER_WG 256
+#endif
+#ifndef LABEL_SMOOTHING
+#define LABEL_SMOOTHING 0.1f
+#endif
+
+constexpr int VEC = 8;
+
+extern "C" __global__ __launch_bounds__(THREADS_PER_WG) void
+fused_ce_loss_fwd(
+    float*                __restrict__ loss_out,  // out: fp32, ROWS
+    float*                __restrict__ max_out,   // out: fp32, ROWS
+    float*                __restrict__ lse_out,   // out: fp32, ROWS
+    const __hip_bfloat16* __restrict__ logits,    // in:  bf16, ROWS*VOCAB
+    const int*            __restrict__ targets)   // in:  int32, ROWS
+{
+  __shared__ float sdata_m[THREADS_PER_WG];
+  __shared__ float sdata_s[THREADS_PER_WG];
+  __shared__ float sdata_sumx[THREADS_PER_WG];
+  __shared__ float sdata_tgt[THREADS_PER_WG];
+
+  const int tid = threadIdx.x;
+  const int row = blockIdx.x;
+  const int target = targets[row];
+  const __hip_bfloat16* row_logits = logits + (size_t)row * VOCAB;
+
+  float m = -INFINITY;
+  float s = 0.0f;
+  float sum_x = 0.0f;
+  float target_logit = 0.0f;
+  constexpr bool needs_sum_x = (LABEL_SMOOTHING != 0.0f);
+
+  // Vectorized stride: each iter loads 8 bf16 = 16 bytes. Warp loads 32*16 = 512 bytes (4 cache lines).
+  const int VOCAB_VEC = VOCAB & ~(VEC - 1);  // round down to multiple of VEC
+  for (int i = tid * VEC; i < VOCAB_VEC; i += THREADS_PER_WG * VEC) {
+    float4 raw = *reinterpret_cast<const float4*>(&row_logits[i]);
+    const __hip_bfloat16* xi = reinterpret_cast<const __hip_bfloat16*>(&raw);
+    #pragma unroll
+    for (int k = 0; k < VEC; k++) {
+      const float x = static_cast<float>(xi[k]);
+      if constexpr (needs_sum_x) sum_x += x;
+      if (i + k == target) target_logit = x;
+      if (x > m) {
+        s = s * __expf(m - x) + 1.0f;
+        m = x;
+      } else {
+        s += __expf(x - m);
+      }
+    }
+  }
+  // tail (VOCAB not divisible by VEC):
+  for (int i = VOCAB_VEC + tid; i < VOCAB; i += THREADS_PER_WG) {
+    const float x = static_cast<float>(row_logits[i]);
+    if constexpr (needs_sum_x) sum_x += x;
+    if (i == target) target_logit = x;
+    if (x > m) { s = s * __expf(m - x) + 1.0f; m = x; }
+    else        { s += __expf(x - m); }
+  }
+
+  sdata_m[tid] = m;
+  sdata_s[tid] = s;
+  sdata_sumx[tid] = sum_x;
+  sdata_tgt[tid] = target_logit;
+  __syncthreads();
+
+  for (int step = THREADS_PER_WG / 2; step > 0; step >>= 1) {
+    if (tid < step) {
+      const float m1 = sdata_m[tid];
+      const float m2 = sdata_m[tid + step];
+      const float s1 = sdata_s[tid];
+      const float s2 = sdata_s[tid + step];
+      const float m_new = fmaxf(m1, m2);
+      const float s_new = s1 * __expf(m1 - m_new) + s2 * __expf(m2 - m_new);
+      sdata_m[tid] = m_new;
+      sdata_s[tid] = s_new;
+      sdata_sumx[tid] += sdata_sumx[tid + step];
+      sdata_tgt[tid]  += sdata_tgt[tid + step];
+    }
+    __syncthreads();
+  }
+
+  if (tid == 0) {
+    const float row_max = sdata_m[0];
+    const float row_sum_exp = sdata_s[0];
+    const float row_sum_x = sdata_sumx[0];
+    const float tgt = sdata_tgt[0];
+    const float row_lse = logf(row_sum_exp) + row_max;
+    const float mean_logits = row_sum_x / static_cast<float>(VOCAB);
+    const float loss = row_lse - (1.0f - LABEL_SMOOTHING) * tgt - LABEL_SMOOTHING * mean_logits;
+    loss_out[row] = loss;
+    max_out[row]  = row_max;
+    lse_out[row]  = row_lse;
+  }
+}

extra/llama_kernels/fused_ce/fused_ce_loss_bwd.cpp

@@ -0,0 +1,58 @@
+#include <hip/hip_runtime.h>
+#include <hip/hip_bf16.h>
+
+// Vectorized CE bwd: 8-wide bf16 loads + stores.
+
+#ifndef VOCAB
+#define VOCAB 128256
+#endif
+#ifndef THREADS_PER_WG
+#define THREADS_PER_WG 256
+#endif
+#ifndef LABEL_SMOOTHING
+#define LABEL_SMOOTHING 0.1f
+#endif
+
+constexpr int VEC = 8;
+
+extern "C" __global__ __launch_bounds__(THREADS_PER_WG) void
+fused_ce_loss_bwd(
+    __hip_bfloat16*       __restrict__ d_logits,
+    const __hip_bfloat16* __restrict__ logits,
+    const float*          __restrict__ lse,
+    const int*            __restrict__ targets,
+    const float*          __restrict__ scale_in)
+{
+  const int tid = threadIdx.x;
+  const int row = blockIdx.x;
+  const int target = targets[row];
+  const float lse_r = lse[row];
+  const __hip_bfloat16* row_logits  = logits   + (size_t)row * VOCAB;
+  __hip_bfloat16*       row_dlogits = d_logits + (size_t)row * VOCAB;
+  const float inv_vocab = 1.0f / static_cast<float>(VOCAB);
+  const float scale = *scale_in;
+  const float ls_term = LABEL_SMOOTHING * inv_vocab;
+
+  const int VOCAB_VEC = VOCAB & ~(VEC - 1);
+  for (int i = tid * VEC; i < VOCAB_VEC; i += THREADS_PER_WG * VEC) {
+    float4 raw = *reinterpret_cast<const float4*>(&row_logits[i]);
+    const __hip_bfloat16* xi = reinterpret_cast<const __hip_bfloat16*>(&raw);
+    __hip_bfloat16 out[VEC];
+    #pragma unroll
+    for (int k = 0; k < VEC; k++) {
+      const float x = static_cast<float>(xi[k]);
+      float g = __expf(x - lse_r);
+      if (i + k == target) g -= (1.0f - LABEL_SMOOTHING);
+      g -= ls_term;
+      out[k] = static_cast<__hip_bfloat16>(g * scale);
+    }
+    *reinterpret_cast<float4*>(&row_dlogits[i]) = *reinterpret_cast<float4*>(out);
+  }
+  for (int i = VOCAB_VEC + tid; i < VOCAB; i += THREADS_PER_WG) {
+    const float x = static_cast<float>(row_logits[i]);
+    float g = __expf(x - lse_r);
+    if (i == target) g -= (1.0f - LABEL_SMOOTHING);
+    g -= ls_term;
+    row_dlogits[i] = static_cast<__hip_bfloat16>(g * scale);
+  }
+}

extra/llama_kernels/fused_mul_quantize_fp8/__init__.py

@@ -0,0 +1,54 @@
+from __future__ import annotations
+import functools, pathlib
+from tinygrad import Tensor, dtypes
+from tinygrad.uop.ops import UOp, Ops, KernelInfo
+from tinygrad.renderer import Estimates
+from extra.llama_kernels import FP8_MAX, NUM_WG, THREADS_PER_WG, compile_cpp, shard_shape, scalar_amax
+
+@functools.cache
+def _custom_mul_quantize_fp8(fp8_out:UOp, amax_buf:UOp, x:UOp, weight:UOp, amax_state:UOp, dname:str) -> UOp:
+  MBS, SEQ, HIDDEN = x.shape
+  n_elems = MBS * SEQ * HIDDEN
+  threads, workgroups = UOp.special(THREADS_PER_WG, "lidx0"), UOp.special(NUM_WG, "gidx0")
+  mem = n_elems * 2 + HIDDEN * 2 + n_elems + NUM_WG * 2
+  sink = UOp.sink(fp8_out.base, amax_buf.base, x.base, weight.base, amax_state.base, threads, workgroups,
+                  arg=KernelInfo(f"fused_mul_quantize_fp8_{n_elems}_h{HIDDEN}", estimates=Estimates(ops=3*n_elems, mem=mem)))
+  src, lib = compile_cpp(pathlib.Path(__file__).parent, "fused_mul_quantize_fp8.cpp", n_elems, HIDDEN)
+  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=(*sink.src, sink)),
+                               UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=lib)))
+
+def _fused_mul_quantize_fp8_bwd(gradient:UOp, kernel:UOp):
+  # NOTE: inputs are (fp8_out, amax_buf, x, weight, amax_state); grads for x and weight
+  _, _, x_u, weight_u, amax_state_u = kernel.src[1:]
+  device = x_u.device
+  grad_t = Tensor(gradient, device=device).cast(dtypes.bfloat16)
+  x_t, weight_t = Tensor(x_u, device=device), Tensor(weight_u, device=device)
+  scale = FP8_MAX / (Tensor(amax_state_u, device=device).float() + 1e-8)
+  grad_scaled = grad_t.float() * scale
+  # NOTE: grad_x stays bf16 to avoid CSE materializing a (MBS, SEQ, HIDDEN) fp32 intermediate
+  grad_x = (grad_scaled * weight_t.float()).cast(dtypes.bfloat16)
+  grad_weight = (grad_scaled * x_t.float()).sum(axis=(0, 1)).cast(dtypes.bfloat16)
+  return (None, None, grad_x.uop, grad_weight.uop, None)
+
+def fused_mul_quantize_fp8(x:Tensor, weight:Tensor, amax_state:Tensor, fp8_dtype) -> tuple[Tensor, Tensor, Tensor]:
+  # NOTE: (x * weight) -> fp8 + amax, delayed scaling. Returns (fp8, inv_scale, new_amax)
+  assert x.dtype == dtypes.bfloat16 and weight.dtype == dtypes.bfloat16
+  assert x.shape[-1] == weight.shape[-1], f"HIDDEN mismatch: x={x.shape}, weight={weight.shape}"
+  MBS, SEQ, HIDDEN = x.shape
+  if isinstance(x.device, tuple):
+    axis, ndev = x.uop.axis, len(x.device)
+    assert axis in (0, 1), f"unsupported sharding axis={axis}"
+    fp8_out = Tensor(Tensor.invalids(*shard_shape((MBS, SEQ, HIDDEN), axis, ndev), dtype=fp8_dtype,
+                                     device=x.device).uop.multi(axis), device=x.device)
+    amax_buf = Tensor(Tensor.invalids(NUM_WG, dtype=dtypes.bfloat16, device=x.device).uop.multi(0), device=x.device)
+    dname = x.device[0].split(":")[0]
+  else:
+    fp8_out = Tensor.invalids(MBS, SEQ, HIDDEN, dtype=fp8_dtype, device=x.device)
+    amax_buf = Tensor.invalids(NUM_WG, dtype=dtypes.bfloat16, device=x.device)
+    dname = x.device.split(":")[0] if isinstance(x.device, str) else x.device
+  fxn = functools.partial(_custom_mul_quantize_fp8, dname=dname)
+  fp8_out, amax_buf, *_ = Tensor.custom_kernel(fp8_out, amax_buf, x, weight, amax_state, fxn=fxn,
+                                                grad_fxn=_fused_mul_quantize_fp8_bwd)
+  new_amax = scalar_amax(amax_buf)
+  inv_scale = (amax_state.float() + 1e-8) / FP8_MAX
+  return fp8_out, inv_scale, new_amax

extra/llama_kernels/rmsnorm/__init__.py

@@ -0,0 +1,24 @@
+from __future__ import annotations
+import functools
+from tinygrad import Tensor
+from tinygrad.uop.ops import UOp
+
+def rmsnorm_fwd(x_in:Tensor, eps:float) -> tuple[Tensor, Tensor]:
+  x = x_in.float()
+  rrms = (x.square().mean(-1, keepdim=True) + eps).rsqrt()
+  return (x * rrms).cast(x_in.dtype), rrms
+
+@functools.cache
+def _rmsnorm_fwd_fxn(x_in_p, eps, device):
+  return rmsnorm_fwd(Tensor(x_in_p, device=device), eps)
+
+def _rmsnorm_bwd(grad:UOp, call:UOp) -> tuple:
+  x_normed = Tensor(call.gettuple(0)).float()
+  do_float = Tensor(grad).float()
+  d_x = Tensor(call.gettuple(1)) * (do_float - x_normed * (do_float * x_normed).mean(-1, keepdim=True))
+  return (d_x.cast(call.src[1].dtype).uop,)
+
+def rmsnorm(x_in:Tensor, eps:float) -> tuple[Tensor, Tensor]:
+  fxn = _rmsnorm_fwd_fxn(x_in.as_param(0).uop, eps, x_in.device)
+  call = UOp.maketuple(fxn[0].uop, fxn[1].uop).call(x_in.uop, grad_fxn=_rmsnorm_bwd)
+  return Tensor(call.gettuple(0)), Tensor(call.gettuple(1))

extra/mmapeak/mmapeak.py

@@ -3,11 +3,12 @@ import os
 # TODO: there is a timing bug without this
 os.environ["AMD_AQL"] = "1"
 
-from tinygrad import Tensor, Device
+from tinygrad import Tensor, Device, GlobalCounters, Context
 from tinygrad.helpers import getenv, DEV
 from tinygrad.uop.ops import UOp, Ops, KernelInfo
 from tinygrad.renderer import Estimates
 from tinygrad.renderer.amd.dsl import Reg, Inst, s, v
+from tinygrad.engine.realize import run_linear
 
 NUM_WORKGROUPS = 96
 WAVE_SIZE = 32
@@ -36,11 +37,17 @@ def launchBenchmark(instruction, vgprIndices, dense=True, accum=False, **kwargs)
     gidx = UOp.special(NUM_WORKGROUPS, "gidx0")
     FLOPs = FLOPS_PER_MATMUL * NUM_WAVES * NUM_WORKGROUPS * INTERNAL_LOOP * INSTRUCTIONS_PER_LOOP
     sink = UOp.sink(A.base, threads, gidx, arg=KernelInfo(inst.op.name.lower(), estimates=Estimates(ops=FLOPs, mem=0)))
-    return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg="AMD"), UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))
+    return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=Device.DEFAULT), UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))
   dummy = Tensor.zeros(1).contiguous().realize()
   out = Tensor.custom_kernel(dummy, fxn=fxn)[0]
-  ei = out.schedule()[-1].lower()
-  elapsed = min([ei.run(wait=True) for _ in range(2)])
+  linear = out.schedule_linear()
+  ets = []
+  with Context(DEBUG=2):
+    for _ in range(2):
+      start = GlobalCounters.time_sum_s
+      run_linear(linear)
+      ets.append(GlobalCounters.time_sum_s - start)
+  elapsed = min(ets)
   FLOPs = FLOPS_PER_MATMUL * NUM_WAVES * NUM_WORKGROUPS * INTERNAL_LOOP * INSTRUCTIONS_PER_LOOP
   print(f"{inst.op_name.lower():<29} : {FLOPs/elapsed/10**12:.2f} T(FL)OPS")
 

extra/optimization/test_beam_search.py

@@ -84,7 +84,7 @@ class TestBeamSearch(unittest.TestCase):
     tc = Device[Device.DEFAULT].renderer.tensor_cores[0]
     size = max(tc.dims[0], tc.dims[1]) * 8
     a, b = Tensor.rand(size, size, dtype=tc.dtype_in), Tensor.rand(size, size, dtype=tc.dtype_in)
-    ast = a.matmul(b, dtype=tc.dtype_out).schedule()[-1].ast
+    ast = a.matmul(b, dtype=tc.dtype_out).schedule_linear().src[-1].src[0]
     s = Scheduler(ast, Device[Device.DEFAULT].renderer)
     s.apply_opt(Opt(OptOps.TC, 0, (-1, 0, 1)))
     up = prod([x for x, t in zip(s.full_shape, s.axis_types) if t in (AxisType.UPCAST, AxisType.UNROLL)])
@@ -94,7 +94,7 @@ class TestBeamSearch(unittest.TestCase):
 
   def test_max_up(self):
     a = Tensor.rand(16, 16)
-    ast = a.schedule()[-1].ast
+    ast = a.schedule_linear().src[-1].src[0]
     s = Scheduler(ast, Device[Device.DEFAULT].renderer)
     for max_up in (2, 4):
       actions = get_kernel_actions(s, include_0=False, max_up=max_up)

test/amd/test_custom_kernel.py

@@ -3,12 +3,14 @@ import functools
 import numpy as np
 from tinygrad import Tensor, Device, dtypes
 from tinygrad.uop.ops import UOp, Ops, KernelInfo
+from tinygrad.engine.realize import run_linear, estimate_uop
 from tinygrad.renderer import Estimates
 from tinygrad.dtype import AddrSpace
+from tinygrad.helpers import getenv
 from tinygrad.runtime.autogen.amd.rdna3.ins import *
 import tinygrad.runtime.autogen.amd.rdna3.ins as r3
 import tinygrad.runtime.autogen.amd.rdna4.ins as r4
-from tinygrad.renderer.amd.dsl import s, v
+from tinygrad.renderer.amd.dsl import s, v, NULL
 from test.amd.helpers import TARGET_TO_ARCH
 from extra.gemm.amd_asm_matmul import Kernel
 
@@ -100,31 +102,44 @@ def custom_lds_sync(A:UOp, arch:str) -> UOp:
 def custom_handwritten(A:UOp, arch:str) -> UOp:
   A = A.flatten()
   threads = UOp.special(128, "lidx0")
-  wg = UOp.special(256, "gidx0")
+  wg = UOp.special(1, "gidx0")
   lds = UOp(Ops.DEFINE_LOCAL, dtypes.uint8.ptr(size=512, addrspace=AddrSpace.LOCAL), (), 'lds')  # 128 * 4 bytes
+  pipes = {getenv("PIPE", "")} if getenv("PIPE", "") else {"SALU", "VALU", "TRANSCENDENTAL", "WMMA"}
   k = Kernel(arch)
-  k.emit(r4.s_nop(0))
-  k.emit(r4.v_mov_b32_e32(v[1], 4))
-  def emit_alt():
-    for i in range(2):
-      k.emit(r4.v_mov_b32_e32(v[20+i], 4.0))
-      k.emit(r4.v_rcp_f32_e32(v[22+i], v[20+i]))
-      k.emit(r4.s_mov_b32(s[20+i], i))
-      k.emit(r4.s_mul_i32(s[14+i], s[12+i], 32))
-  def emit_wmma():
-    for _ in range(2):
-      k.emit(r4.v_wmma_f32_16x16x16_f16(v[0:7], v[8:11], v[8:11], 1))
-  k.label("start")
-  k.emit(s_mov_b32(s[1], 10))
+  # wrap in loop to filter out icache misses
+  LOOP_N, UNROLL_N = 8, 5
+  k.emit(r4.s_mov_b32(s[1], LOOP_N))
   k.label("loop")
-  # wmma should've overlapped here if it was a different unit?
-  for _ in range(2):
-    emit_wmma()
-    emit_alt()
-  for _ in range(8): k.emit(s_nop(1))
-  k.emit(s_add_u32(s[1], s[1], -1))
-  k.emit(s_cmp_eq_i32(s[1], 0))
-  k.emit(s_cbranch_scc0(), target="loop")
+  if "SALU" in pipes:
+    for i in range(UNROLL_N):
+      k.emit(r4.s_mov_b32(s[20+i], i))
+      k.emit(r4.s_min_i32(s[30+i], i))
+      k.emit(r4.s_mov_b32(s[40+i], i))
+      k.emit(r4.s_mul_i32(s[14+i], s[12+i], 32))
+  if "VALU" in pipes:
+    for i in range(UNROLL_N):
+      k.emit(r4.v_mov_b32_e32(v[20+i], i))
+      k.emit(r4.v_lshlrev_b64_e32(v[30+2*i:31+2*i], 2, v[12+i:13+i]))
+      k.emit(r4.v_mad_co_u64_u32(v[40+2*i:41+2*i], NULL, v[12+i], v[13+i], v[14+i:15+i]))
+  if "TRANSCENDENTAL" in pipes:
+    # transcendental VALU runs on the TFU, it can run regular VALU at the same time
+    for i in range(UNROLL_N):
+      k.emit(r4.v_mov_b32_e32(v[20+i], i))
+      k.emit(r4.v_s_rcp_f32(s[20+i], s[12+i]))
+      k.emit(r4.v_rcp_f32_e32(v[30+i], v[12+i]))
+      k.emit(r4.v_s_exp_f32(s[30+i], s[12+i]))
+  if "WMMA" in pipes:
+    base = 30
+    for i in range(UNROLL_N):
+      a = base + i*40
+      b, cd = a + 4, a + 8
+      k.emit(r4.v_wmma_f32_16x16x16_f16(v[cd:cd+7], v[a:a+3], v[b:b+3], v[cd:cd+7]))
+      a = base + i*40 + 16
+      b, cd = a + 2, a + 4
+      k.emit(r4.v_wmma_i32_16x16x16_iu8(v[cd:cd+7], v[a:a+1], v[b:b+1], v[cd:cd+7]))
+  k.emit(r4.s_add_co_i32(s[1], s[1], -1))
+  k.emit(r4.s_cmp_eq_i32(s[1], 0))
+  k.emit(r4.s_cbranch_scc0(), target="loop")
   k.emit(r4.s_endpgm())
   insts = k.finalize()
   sink = UOp.sink(A.base, threads, wg, lds, arg=KernelInfo("custom_handwritten"))
@@ -154,10 +169,11 @@ class TestCustomKernel(unittest.TestCase):
     if self.arch != "rdna3": self.skipTest("only rdna3")
     a = Tensor.full((16, 16), 1.).contiguous().realize()
     a = Tensor.custom_kernel(a, fxn=custom_add_one)[0]
-    ei = a.schedule()[-1].lower()
-    self.assertEqual(ei.prg.estimates.ops, a.numel())
-    self.assertEqual(ei.prg.estimates.mem, a.nbytes()*2)
-    ei.run()
+    linear = a.schedule_linear()
+    est = estimate_uop(linear.src[-1])
+    self.assertEqual(est.ops, a.numel())
+    self.assertEqual(est.mem, a.nbytes()*2)
+    run_linear(linear)
     self.assertTrue((a.numpy() == 2.).all())
 
   def test_variable(self):
@@ -165,9 +181,9 @@ class TestCustomKernel(unittest.TestCase):
     b = Tensor.full((16, 16), 1, dtype=dtypes.uint32).contiguous().realize()
     a = Tensor.zeros_like(b).contiguous().realize()
     a = Tensor.custom_kernel(a, b, fxn=custom_add_var)[0]
-    ei = a.schedule()[-1].lower()
+    linear = a.schedule_linear()
     for i in range(4):
-      ei.run({"var":i})
+      run_linear(linear, var_vals={"var":i})
       self.assertTrue((a.numpy() == 1+i).all())
 
   def test_lds_sync(self):

test/amd/test_integration.py

@@ -78,18 +78,18 @@ class TestTinygradIntegration(unittest.TestCase):
   def _get_kernel_code(self, op_fn) -> bytes:
     from tinygrad import Tensor
     from tinygrad.helpers import Target
-    from tinygrad.codegen import get_program
+    from tinygrad.codegen import to_program
     from tinygrad.renderer.llvmir import AMDLLVMRenderer
     from tinygrad.runtime.support.elf import elf_loader
     from tinygrad.uop.ops import Ops
 
     result = op_fn(Tensor)
-    schedule = result.schedule()
-    sink_items = [si for si in schedule if si.ast.op == Ops.SINK]
+    linear = result.schedule_linear()
+    sink_items = [call for call in linear.src if call.src[0].op == Ops.SINK]
     assert len(sink_items) > 0, "No SINK in schedule"
     renderer = AMDLLVMRenderer(Target("AMD", arch='gfx1100'))
-    prg = get_program(sink_items[0].ast, renderer)
-    lib = renderer.compiler.compile(prg.src)
+    prg = to_program(sink_items[0].src[0], renderer)
+    lib = renderer.compiler.compile(prg.src[3].arg)
     return next(s.content for s in elf_loader(lib)[1] if s.name == ".text")
 
   def test_simple_add_kernel(self):

test/amd/test_mockgpu_invalid.py

@@ -14,8 +14,8 @@ from tinygrad.runtime.ops_amd import AMDProgram
 dev = Device["AMD"]
 a = Tensor([1.0]).realize()
 b = a + 1
-si = b.schedule()[-1]
-runner = get_runner(dev.device, si.ast)
+si = b.schedule_linear().src[-1]
+runner = get_runner(dev.device, si.src[0])
 
 prg = runner._prg
 lib = bytearray(prg.lib)

test/amd/test_roundtrip.py

@@ -57,49 +57,53 @@ class KernelSnapshot:
 def get_kernels_from_tinygrad(op_fn) -> tuple[list[KernelSnapshot], dict[int, int], dict[int, bytes]]:
   """Compile a tinygrad operation and extract all kernels with their buffer mappings."""
   from tinygrad import Tensor
+  from tinygrad.uop.ops import Ops
+  from tinygrad.engine.realize import compile_linear, resolve_params, unwrap_multi
   from tinygrad.runtime.support.elf import elf_loader
 
   out = op_fn(Tensor)
-  sched = out.schedule()
+  linear = compile_linear(out.schedule_linear())
   kernels = []
   buf_pool: dict[int, int] = {}  # buffer id -> size
   buf_data: dict[int, bytes] = {}  # buffer id -> initial data from COPY
 
-  for ei in sched:
-    lowered = ei.lower()
-    if ei.ast.op.name == 'COPY':
-      # Handle COPY: extract source data to initialize destination buffer
-      if len(lowered.bufs) >= 2:
-        dst_buf, src_buf = lowered.bufs[0], lowered.bufs[1]
-        dst_id = id(dst_buf)
-        if dst_id not in buf_pool:
-          buf_pool[dst_id] = dst_buf.nbytes
-        # Get source data if it's from numpy/CPU
-        if hasattr(src_buf, 'base') and src_buf.base is not None and hasattr(src_buf.base, '_buf'):
-          src_data = bytes(src_buf.base._buf)
-          buf_data[dst_id] = src_data
-    elif ei.ast.op.name == 'SINK':
-      if lowered.prg and lowered.prg.p.lib:
-        lib = bytes(lowered.prg.p.lib)
-        _, sections, _ = elf_loader(lib)
-        for sec in sections:
-          if sec.name == '.text':
-            buf_idxs = []
-            buf_sizes = []
-            for b in lowered.bufs:
-              buf_id = id(b)
-              if buf_id not in buf_pool:
-                buf_pool[buf_id] = b.nbytes
-              buf_idxs.append(buf_id)
-              buf_sizes.append(b.nbytes)
-            kernels.append(KernelSnapshot(
-              code=bytes(sec.content),
-              src=lowered.prg.p.src,
-              global_size=tuple(lowered.prg.p.global_size),
-              local_size=tuple(lowered.prg.p.local_size),
-              buf_idxs=buf_idxs,
-              buf_sizes=buf_sizes
-            ))
+  for call in linear.src:
+    ast = call.src[0]
+    for bufs, _ in unwrap_multi(call, resolve_params(call, ())):
+      if ast.op is Ops.COPY:
+        # Handle COPY: extract source data to initialize destination buffer
+        if len(bufs) >= 2:
+          dst_buf, src_buf = bufs[0], bufs[1]
+          dst_id = id(dst_buf)
+          if dst_id not in buf_pool:
+            buf_pool[dst_id] = dst_buf.nbytes
+          # Get source data if it's from numpy/CPU
+          if hasattr(src_buf, 'base') and src_buf.base is not None and hasattr(src_buf.base, '_buf'):
+            src_data = bytes(src_buf.base._buf)
+            buf_data[dst_id] = src_data
+      elif ast.op is Ops.PROGRAM:
+        info = ast.arg
+        if len(ast.src) > 4 and ast.src[4].op is Ops.BINARY:
+          lib = bytes(ast.src[4].arg)
+          _, sections, _ = elf_loader(lib)
+          for sec in sections:
+            if sec.name == '.text':
+              buf_idxs = []
+              buf_sizes = []
+              for b in bufs:
+                buf_id = id(b)
+                if buf_id not in buf_pool:
+                  buf_pool[buf_id] = b.nbytes
+                buf_idxs.append(buf_id)
+                buf_sizes.append(b.nbytes)
+              kernels.append(KernelSnapshot(
+                code=bytes(sec.content),
+                src=ast.src[3].arg,
+                global_size=tuple(info.global_size),
+                local_size=tuple(info.local_size),
+                buf_idxs=buf_idxs,
+                buf_sizes=buf_sizes
+              ))
   if not kernels: raise RuntimeError("No kernel found")
   return kernels, buf_pool, buf_data
 

test/amd/test_sqtt_examples.py

@@ -8,10 +8,11 @@ from tinygrad.runtime.support.elf import elf_loader
 from tinygrad.renderer.amd import decode_inst
 from tinygrad.runtime.autogen.amd.rdna3.ins import SOPP
 from tinygrad.runtime.autogen.amd.rdna3.enum import SOPPOp
-from tinygrad.renderer.amd.sqtt import (decode, LAYOUT_HEADER, WAVESTART, WAVESTART_RDNA4, WAVEEND, INST, INST_RDNA4, VALUINST,
+from tinygrad.renderer.amd.sqtt import (decode, LAYOUT_HEADER, WAVESTART, WAVESTART_RDNA4, WAVEEND, WAVEEND_RDNA4, INST, INST_RDNA4, VALUINST,
                                      IMMEDIATE, IMMEDIATE_MASK, PACKET_TYPES_RDNA3, PACKET_TYPES_RDNA4, PACKET_TYPES_CDNA, CDNA_WAVESTART,
                                      print_packets, CDNA_WAVEEND, CDNA_INST)
 from test.amd.helpers import TARGET_TO_ARCH
+from test.amd.test_sqttmap import needs_rocprof
 
 import tinygrad
 EXAMPLES_DIR = Path(tinygrad.__file__).parent.parent / "extra/sqtt/examples"
@@ -132,7 +133,7 @@ class SQTTExamplesTestBase(unittest.TestCase):
       with self.subTest(example=name):
         all_packets = [p for e in events for p in decode(e.blob)]
         self.assertGreater(len([p for p in all_packets if isinstance(p, (WAVESTART, WAVESTART_RDNA4, CDNA_WAVESTART))]), 0, f"no WAVESTART in {name}")
-        self.assertGreater(len([p for p in all_packets if isinstance(p, (WAVEEND, CDNA_WAVEEND))]), 0, f"no WAVEEND in {name}")
+        self.assertGreater(len([p for p in all_packets if isinstance(p, (WAVEEND, WAVEEND_RDNA4, CDNA_WAVEEND))]), 0, f"no WAVEEND in {name}")
 
   def test_time_monotonic(self):
     for name, (events, *_) in self.examples.items():
@@ -160,6 +161,7 @@ class SQTTExamplesTestBase(unittest.TestCase):
         counts = [len(list(decode(e.blob))) for e in events]
         self.assertEqual(counts, self.expected[name], f"packet count mismatch in {name}")
 
+  @needs_rocprof
   def test_rocprof_wave_times_match(self):
     """Wave start/end times must match rocprof exactly."""
     for name, (events, lib, base) in self.examples.items():
@@ -180,7 +182,7 @@ class SQTTExamplesTestBase(unittest.TestCase):
           for p in decode(event.blob):
             if first_timestamp is None: first_timestamp = p._time
             if isinstance(p, (WAVESTART, CDNA_WAVESTART, WAVESTART_RDNA4)): wave_starts[(p.wave, p.simd, p.cu)] = p._time
-            elif isinstance(p, (WAVEEND, CDNA_WAVEEND)) and (key := (p.wave, p.simd, p.cu)) in wave_starts:
+            elif isinstance(p, (WAVEEND, WAVEEND_RDNA4, CDNA_WAVEEND)) and (key := (p.wave, p.simd, p.cu)) in wave_starts:
               our_waves.append((wave_starts[key], p._time))
           for st in wave_starts.values():
             self.assertGreater(st, first_timestamp, "wave start must be after the first packet")
@@ -189,6 +191,7 @@ class SQTTExamplesTestBase(unittest.TestCase):
         for st, et in our_waves:
           self.assertGreater(et, st, "wave end must be after start")
 
+  @needs_rocprof
   def test_rocprof_inst_times_match(self):
     """Instruction times must match rocprof exactly (excluding s_endpgm)."""
     for name, (events, lib, base) in self.examples.items():

test/amd/test_sqtt_profiler.py

@@ -1,6 +1,8 @@
 import unittest, contextlib
 from tinygrad import Device, Tensor, Context, TinyJit
 from tinygrad.device import Compiled, ProfileProgramEvent, ProfileDeviceEvent
+from tinygrad.engine.realize import run_linear
+from tinygrad.codegen import to_program
 from tinygrad.viz.serve import load_amd_counters, VizData
 
 @contextlib.contextmanager
@@ -26,39 +28,41 @@ class TestSQTTProfiler(unittest.TestCase):
   def test_simple(self):
     t = Tensor.empty(1) + 1
     with save_sqtt() as sqtt:
-      ei = t.schedule()[0].lower()
-      ei.run()
+      linear = t.schedule_linear()
+      run_linear(linear)
+    fn_name = to_program(linear.src[0].src[0], renderer=Device[Device.DEFAULT].renderer).arg.function_name
     self.assertEqual(len(sqtt), 1)
-    self.assertEqual(sqtt[0]["name"], f"SQTT {ei.prg.p.function_name}")
+    self.assertEqual(sqtt[0]["name"], f"SQTT {fn_name}")
 
   def test_multiple_runs(self):
     t = Tensor.empty(1) + 1
     with save_sqtt() as sqtt:
-      ei = t.schedule()[0].lower()
-      for _ in range(N:=3):
-        ei.run()
+      linear = t.schedule_linear()
+      for _ in range(N:=3): run_linear(linear)
+    fn_name = to_program(linear.src[0].src[0], renderer=Device[Device.DEFAULT].renderer).arg.function_name
     self.assertEqual(len(sqtt), N)
     for i in range(1, N):
-      self.assertEqual(sqtt[i]["name"], f"SQTT {ei.prg.p.function_name} n{i+1}")
+      self.assertEqual(sqtt[i]["name"], f"SQTT {fn_name} n{i+1}")
 
   def test_multiple_kernels(self):
     t = ((Tensor.empty(1) + 1).contiguous() + 2)
-    sched = t.schedule()
+    linear = t.schedule_linear()
     with save_sqtt() as sqtt:
-      for si in sched: si.lower().run()
-    self.assertEqual(len(sqtt), len(sched))
-    for i,k in enumerate(sched):
-      self.assertEqual(sqtt[i]["name"], f"SQTT {k.lower().prg.p.function_name}")
+      run_linear(linear)
+    self.assertEqual(len(sqtt), len(linear.src))
+    for i,call in enumerate(linear.src):
+      fn_name = to_program(call.src[0], renderer=Device[Device.DEFAULT].renderer).arg.function_name
+      self.assertEqual(sqtt[i]["name"], f"SQTT {fn_name}")
 
   def test_multiple_kernels_lower(self):
     t = ((Tensor.empty(1) + 1).contiguous() + 2)
-    sched = t.schedule()
+    linear = t.schedule_linear()
     with save_sqtt() as sqtt:
-      prgs = [si.lower() for si in sched]
-      for p in prgs: p.run()
-    self.assertEqual(len(sqtt), len(sched))
-    for i,ei in enumerate(prgs):
-      self.assertEqual(sqtt[i]["name"], f"SQTT {ei.prg.p.function_name}")
+      run_linear(linear)
+    self.assertEqual(len(sqtt), len(linear.src))
+    for i,call in enumerate(linear.src):
+      fn_name = to_program(call.src[0], renderer=Device[Device.DEFAULT].renderer).arg.function_name
+      self.assertEqual(sqtt[i]["name"], f"SQTT {fn_name}")
 
   def test_jit(self):
     @TinyJit

test/amd/test_sqttmap.py

@@ -1,19 +1,36 @@
 # test to compare every packet with the rocprof decoder
-import unittest, pickle
+import unittest, pickle, functools
 from typing import Iterator
 from pathlib import Path
-from tinygrad.helpers import DEBUG, getenv, temp, ansistrip
+from tinygrad.helpers import DEBUG, getenv, temp, ansistrip, Context
 from tinygrad.renderer.amd.sqtt import print_packets, map_insts
 from tinygrad.runtime.autogen.amd.rdna3.ins import s_endpgm
-from tinygrad.viz.serve import sqtt_timeline
+from tinygrad.viz.serve import sqtt_timeline, amd_decode
 from test.amd.disasm import disasm
 from test.null.test_viz import run_cli
 
 import tinygrad
 EXAMPLES_DIR = Path(tinygrad.__file__).parent.parent / "extra/sqtt/examples"
 
+def needs_rocprof(fn):
+  @functools.wraps(fn)
+  def wrapper(self, *args, **kwargs):
+    # check if latest rocprof is available, if not, skip rocprof comparison tests
+    # rocprof doesn't have a version string, decode a known pickle to validate it's the latest
+    try:
+      from extra.sqtt.roc import decode as roc_decode
+      with open(EXAMPLES_DIR/"gfx1200"/"profile_plus_run_0.pkl", "rb") as f:
+        data = pickle.load(f)
+      sqtt = [e for e in data if type(e).__name__ == "ProfileSQTTEvent"][1]
+      kern = {e.tag:e for e in data if type(e).__name__ == "ProfileProgramEvent"}[sqtt.kern]
+      rctx = roc_decode([sqtt], {kern.tag:{addr+kern.base:inst for addr,inst in amd_decode(kern.lib, "gfx1200").items()}})
+      insts = [e.time for e in list(rctx.inst_execs.values())[0][0].unpack_insts()]
+      self.assertListEqual(insts, [28178, 28179, 28180, 28181, 28182, 29882, 29883, 29884, 29885, 30966, 30983, 30985, 30992, 30993])
+    except Exception as e: self.skipTest(f"latest rocprof not available, install with extra/sqtt/install_rocprof_decoder.py: {e}")
+    return fn(self, *args, **kwargs)
+  return wrapper
+
 def rocprof_inst_traces_match(sqtt, prg, target):
-  from tinygrad.viz.serve import amd_decode
   from extra.sqtt.roc import decode as roc_decode, InstExec
   addr_table = amd_decode(prg.lib, target)
   disasm_map = {addr+prg.base:inst for addr,inst in addr_table.items()}
@@ -63,6 +80,7 @@ class TestSQTTMapBase(unittest.TestCase):
       if sqtt_events and kern_events:
         cls.examples[pkl_path.stem] = (sqtt_events, kern_events, cls.target)
 
+  @needs_rocprof
   def test_rocprof_inst_traces_match(self):
     for name, (events, kern_events, target) in self.examples.items():
       if "sync" in name and self.target.startswith("gfx12"):
@@ -94,7 +112,7 @@ class TestSQTTMapBase(unittest.TestCase):
             elif "WAVE" in e.device:
               # sopk/immediates don't get ALU/MEM EXEC
               if e.name.display_name not in {"IMMEDIATE", "IMMEDIATE_MASK", "JUMP", "JUMP_NO", "MESSAGE", "BARRIER", "BARRIER_SIGNAL",
-                                             "WAVEEND", "WAVERDY"} and not e.name.display_name.startswith("OTHER_"): insts += 1
+                                             "WAVEEND", "WAVEEND_RDNA4", "WAVERDY"} and not e.name.display_name.startswith("OTHER_"): insts += 1
             else: raise Exception(f"timeline row must be INST or EXEC, got {e.device}")
           self.assertEqual(execs, insts)
 
@@ -111,15 +129,18 @@ class TestSQTTMapBase(unittest.TestCase):
 
   def test_sqtt_cli(self):
     for pkl_path in sorted((EXAMPLES_DIR/self.target).glob("*.pkl")):
-      out = run_cli("--profile", "--profile-path", str(pkl_path))
+      out = run_cli("--profile-path", str(pkl_path), "--ls")
       sqtt_traces = [l.strip() for l in out.split("\n") if "SQTT" in l]
       for name in sqtt_traces:
-        out = run_cli("--profile", "--profile-path", str(pkl_path), "-s", ansistrip(name))
+        out = run_cli("--profile-path", str(pkl_path), "-s", ansistrip(name))
         lines = out.split("\n")
         self.assertIn("Clk", lines[0])
         for r in lines[2:]:
           parts = r.split()
           self.assertTrue(parts[0].isdigit(), f"expected clock timestamp, got {parts[0]}")
+      with Context(DEBUG=2):
+        kernels = run_cli("--profile-path", str(pkl_path), "-s", "AMD").split("\n")
+      self.assertEqual(len(kernels), len(self.examples[pkl_path.stem][1]))
 
 class TestSQTTMapRDNA3(TestSQTTMapBase): target = "gfx1100"
 
@@ -127,14 +148,25 @@ class TestSQTTMapRDNA4(TestSQTTMapBase):
   target = "gfx1200"
 
   @unittest.expectedFailure
-  def test_rdna4_wmma(self):
+  def test_pipes(self):
     events, kernels, target = self.examples["profile_handwritten_run_0"]
-    row_ends = {}
-    for e in sqtt_timeline(events[0].blob, list(kernels.values())[0].lib, target):
-      if type(e).__name__ != "ProfileRangeEvent" or e.device != "ALUEXEC:0 WMMA": continue
-      if (et:=row_ends.get(e.device)) is not None and e.st < et:
-        raise RuntimeError(f"WMMA exec overlaps in {e.device}: {e.st} {et}.")
-      row_ends[e.device] = e.en
+    lib = list(kernels.values())[0].lib
+    dispatch_st:dict[str, int] = {}
+    row_ends:dict[str, int] = {}
+    row_counts:dict[str, int] = {}
+    for e in sqtt_timeline(events[1].blob, lib, target):
+      if type(e).__name__ != "ProfileRangeEvent": continue
+      info = e.name.ret or ""
+      if e.device.startswith("WAVE"):
+        idx = row_counts.get(e.device, 0)
+        dispatch_st[f"{e.device}-{idx}"] = int(e.st)
+        row_counts[e.device] = idx + 1
+      elif info.startswith("LINK:"):
+        delay = int(e.st) - dispatch_st[info[len("LINK:"):]]
+        self.assertGreaterEqual(delay, 1, f"EXEC {e.device} starts before DISPATCH: delay={delay}")
+        if (prev_en:=row_ends.get(e.device)) is not None:
+          self.assertGreaterEqual(e.st, prev_en, f"EXEC overlap in {e.device}: {e.st} < prev end {prev_en}")
+        row_ends[e.device] = int(e.en)
 
 class TestSQTTMapCDNA(TestSQTTMapBase):
   target = "gfx950"

test/backend/test_arange.py

@@ -2,23 +2,18 @@ import unittest
 import numpy as np
 from tinygrad import Tensor, GlobalCounters, dtypes, nn, Device, Variable
 from tinygrad.helpers import Context, getenv, DEV
-from tinygrad.engine.realize import run_linear
-from tinygrad.schedule import linear_to_schedule
-from tinygrad.engine.realize import CompiledRunner, get_program
-from tinygrad.schedule import ExecItem
-from tinygrad.renderer import Estimates
+from tinygrad.engine.realize import run_linear, estimate_uop
 from tinygrad.renderer.ptx import PTXRenderer
 from test.helpers import needs_second_gpu
 
 class TestArange(unittest.TestCase):
   def _get_flops(self, tensor, desired):
     GlobalCounters.reset()
-    sched = tensor.schedule()
-    self.assertEqual(len(sched), 1)
-    p = get_program(sched[-1].ast, renderer=Device[Device.DEFAULT].renderer)
-    ExecItem(sched[-1].ast, [tensor.uop.buffer], prg=CompiledRunner(p)).run()
+    linear = tensor.schedule_linear()
+    self.assertEqual(len(linear.src), 1)
+    run_linear(linear)
     np.testing.assert_equal(tensor.numpy(), desired)
-    return p.estimates.ops
+    return estimate_uop(linear.src[-1]).ops
 
   def test_arange_complexity(self):
     self.assertEqual(self._get_flops(Tensor.arange(256), np.arange(256)), 0)
@@ -41,9 +36,8 @@ class TestArange(unittest.TestCase):
   def test_tri_complexity(self):
     with Context(NOOPT=1):
       t = Tensor.ones(256, 256).contiguous().realize()
-      sched = t.triu().schedule()
-      p = get_program(sched[-1].ast, renderer=Device[Device.DEFAULT].renderer)
-      self.assertLessEqual(Estimates.from_uops(p.uops).ops, 4 * 256 * 256)
+      linear = t.triu().schedule_linear()
+      self.assertLessEqual(estimate_uop(linear.src[-1]).ops, 4 * 256 * 256)
 
 DSET, DDIM = 2048, 32
 
@@ -56,7 +50,7 @@ class TestIndexing(unittest.TestCase):
       GlobalCounters.reset()
       out = ((Tensor.arange(1,16385)-1)*needle).sum()
       linear, var_vals = out.linear_with_vars()
-      self.assertEqual(len(linear_to_schedule(linear)), 1)
+      self.assertEqual(len(linear.src), 1)
       run_linear(linear, var_vals)
     self.assertEqual(out.item(), 1337)
 
@@ -73,7 +67,7 @@ class TestIndexing(unittest.TestCase):
       full = (rng==idxs).where(reshape_dataset, Tensor.zeros(4, DDIM, DSET, 1))
       X = full.sum(axis=(2,3))
       linear, var_vals = X.linear_with_vars()
-      self.assertEqual(len(linear_to_schedule(linear)), 1)
+      self.assertEqual(len(linear.src), 1)
       run_linear(linear, var_vals)
       assert GlobalCounters.global_ops < 4*DSET, f"too many ops {GlobalCounters.global_ops}"
     np.testing.assert_allclose(real_index, X.numpy())
@@ -99,7 +93,7 @@ class TestIndexing(unittest.TestCase):
       X = dataset[idxs]
       assert X.shape == (4,DDIM)
       linear, var_vals = X.linear_with_vars()
-      self.assertEqual(len(linear_to_schedule(linear)), 1)
+      self.assertEqual(len(linear.src), 1)
       run_linear(linear, var_vals)
       assert GlobalCounters.global_ops < 4*DSET, f"too many ops {GlobalCounters.global_ops}"
     np.testing.assert_allclose(real_index, X.numpy())
@@ -114,7 +108,7 @@ class TestIndexing(unittest.TestCase):
       X = dataset[idxs]
       assert X.shape == (4,DDIM)
       linear, var_vals = X.linear_with_vars()
-      self.assertEqual(len(linear_to_schedule(linear)), 1)
+      self.assertEqual(len(linear.src), 1)
       run_linear(linear, var_vals)
       assert GlobalCounters.global_ops < 4*DSET, f"too many ops {GlobalCounters.global_ops} != {4*DSET}"
     np.testing.assert_allclose(real_index, X.numpy())
@@ -235,10 +229,9 @@ class TestIndexing(unittest.TestCase):
     xq = xq.reshape(bs, seqlen, n_heads, head_dim)
     xq_rope, _ = apply_rotary_emb(xq, xq, freqs_cis)
     xq_rope.sum().backward()
-    sched = wq.grad.schedule()
-    assert len(sched) == 1, f"expected one kernel for backward, got: {len(sched)}"
-    prg = sched[0].lower().prg.p
-    bwd_ops = prg.estimates.ops
+    linear = wq.grad.schedule_linear()
+    assert len(linear.src) == 1, f"expected one kernel for backward, got: {len(linear.src)}"
+    bwd_ops = estimate_uop(linear.src[0]).ops
     # bfloat16 on non CDNA4 has ~10x ops overhead because of the software emulation
     if dtype == dtypes.bfloat16 and not Device[Device.DEFAULT].renderer.target.arch.startswith("gfx950"): ops_scale = 10
     else: ops_scale = 1

test/backend/test_asm_gemm.py

@@ -46,9 +46,9 @@ def run_asm_gemm(a_shape, b_shape, dtype=dtypes.float16, a_shard=None, b_shard=N
       np.testing.assert_allclose(tst.numpy(), ref.numpy(), atol=atol, rtol=rtol)
       np.testing.assert_allclose(a.grad.numpy(), a_ref.grad.numpy(), atol=grad_atol, rtol=grad_rtol)
       np.testing.assert_allclose(b.grad.numpy(), b_ref.grad.numpy(), atol=grad_atol, rtol=grad_rtol)
-    assert tst.allclose(ref, atol=atol, rtol=rtol), "forward mismatch"
-    assert a.grad.allclose(a_ref.grad, atol=grad_atol, rtol=grad_rtol), "grad_a mismatch"
-    assert b.grad.allclose(b_ref.grad, atol=grad_atol, rtol=grad_rtol), "grad_b mismatch"
+    assert tst.allclose(ref, atol=atol, rtol=rtol).item(), "forward mismatch"
+    assert a.grad.allclose(a_ref.grad, atol=grad_atol, rtol=grad_rtol).item(), "grad_a mismatch"
+    assert b.grad.allclose(b_ref.grad, atol=grad_atol, rtol=grad_rtol).item(), "grad_b mismatch"
 
 def verify_asm_gemm(batch:int, M:int, N:int, K:int, dtype=dtypes.float16, gpus:int=1) -> None:
   run_asm_gemm((batch, M, K), (K, N), dtype=dtype, a_shard=0, b_shard=None, gpus=gpus)

test/backend/test_const_folding.py

@@ -8,8 +8,8 @@ from test.helpers import not_support_multi_device
 
 def _check_ast_count(desired_count:int, t:Tensor):
   # NOTE: this has side effect because everything can be scheduled only once
-  schedule = t.schedule()
-  asts = [s for s in schedule if s.ast.op is Ops.SINK]
+  schedule = t.schedule_linear()
+  asts = [s for s in schedule.src if s.src[0].op is Ops.SINK]
   len(asts)
   # NOT SUPPORTED ANYMORE
   #assert len(asts) == desired_count, f"{len(asts)} != {desired_count}"
@@ -28,8 +28,8 @@ class TestMovedConstFolding(unittest.TestCase):
     _check_ast_count(1, Tensor([1.0, 2, 3, 4]) * Tensor.ones(2).pad(((1, 1),)))
 
   def test_copy_padded_const(self):
-    schedule = Tensor.ones(4, device="CPU:0").pad(((1, 1),)).to("CPU:1").schedule()
-    assert not any(si.ast.op is Ops.COPY for si in schedule), "const copy should be folded"
+    schedule = Tensor.ones(4, device="CPU:0").pad(((1, 1),)).to("CPU:1").schedule_linear()
+    assert not any(si.src[0].op is Ops.COPY for si in schedule.src), "const copy should be folded"
     np.testing.assert_equal(Tensor.ones(4, device="CPU:0").pad(((1, 1),)).to("CPU:1").numpy(), [0, 1, 1, 1, 1, 0])
 
   def test_cast_padded(self):

test/backend/test_custom_kernel.py

@@ -189,7 +189,7 @@ class TestCustomKernel(unittest.TestCase):
     A = Tensor.randn(16, 16).contiguous()
     B = Tensor.empty(16)
     B = Tensor.custom_kernel(B, A, fxn=slice_sum_kernel)[0]
-    self.assertTrue(B.allclose(A.sum(1)))
+    self.assertTrue(B.allclose(A.sum(1)).item())
 
   def test_gemm(self):
     N = 16
@@ -273,12 +273,12 @@ class TestCustomKernel(unittest.TestCase):
     C, D, _, _ = Tensor.custom_kernel(C, D, A2, B2, fxn=custom_elementwise_addmul_kernel)  # depends on A2 AND B2
     E = (A2 * 3).contiguous()                      # kernel 2: depends only on A2
     result = (C + D + E).sum()                     # kernel 3: custom_addmul, then kernel 4: sum
-    schedule = result.schedule()
+    schedule = result.schedule_linear().src
 
     # Find the custom_addmul kernel position
     custom_idx = next((i for i, item in enumerate(schedule)
-                       if hasattr(item.ast, "arg") and hasattr(item.ast.arg, "name")
-                       and "custom_addmul" in item.ast.arg.name), None)
+                       if hasattr(item.src[0], "arg") and hasattr(item.src[0].arg, "name")
+                       and "custom_addmul" in item.src[0].arg.name), None)
 
     self.assertIsNotNone(custom_idx, "custom_addmul kernel not found in schedule")
     self.assertEqual(custom_idx, 3, f"custom_addmul should be at index 3, got {custom_idx}")

test/backend/test_graph.py

@@ -6,7 +6,7 @@ from tinygrad.tensor import Tensor
 from tinygrad.helpers import Context, from_mv
 from tinygrad.dtype import dtypes
 from tinygrad.engine.jit import MultiGraphRunner
-from tinygrad.engine.realize import run_linear
+from tinygrad.engine.realize import run_linear, compile_linear
 from tinygrad.uop.ops import UOp, Ops, buffers
 
 from test.helpers import needs_second_gpu
@@ -24,7 +24,7 @@ def get_ast(device:str, num_inputs:int) -> UOp:
       fst = [Tensor.randn(BUF_SIZE, dtype=dtypes.int).realize() for _ in range(num_inputs)]
       s = fst[0]
       for i in range(1, num_inputs): s = s.bitwise_xor(fst[i])
-      cached_asts[(device, num_inputs)] = s.schedule()[-1].ast
+      cached_asts[(device, num_inputs)] = s.schedule_linear().src[-1].src[0]
   return cached_asts[(device, num_inputs)]
 
 def make_buffer(device, size=BUF_SIZE, fill=False):
@@ -44,7 +44,7 @@ def get_buf_uop(buf:Buffer, cache:dict[Buffer,UOp]) -> UOp:
   return cache[buf]
 
 def make_graph(graph_cls, calls:list[UOp]):
-  linear = UOp(Ops.LINEAR, src=tuple(calls))
+  linear = compile_linear(UOp(Ops.LINEAR, src=tuple(calls)))
   cf = UOp(Ops.CUSTOM_FUNCTION, dtypes.void, src=(linear,), arg="graph")
   return graph_cls(cf, [])
 

test/backend/test_jit.py

@@ -92,6 +92,20 @@ class TestJit(unittest.TestCase):
       np.testing.assert_allclose(e.numpy(), a.numpy()*b.numpy(), atol=1e-4, rtol=1e-5)
     assert_jit_cache_len(f, 3)
 
+  def test_global_counters_jit(self):
+    @TinyJit
+    def f(a, b):
+      c = (a + b).realize()
+      d = (c * 2).realize()
+      return (d - a).realize()
+    a, b = Tensor.randn(64, 64).realize(), Tensor.randn(64, 64).realize()
+    for _ in range(4):
+      GlobalCounters.reset()
+      f(a, b)
+      Device[a.device].synchronize()
+      self.assertGreater(GlobalCounters.global_mem, 0)
+      self.assertGreater(GlobalCounters.global_ops, 0)
+
   def test_nothing_jitted(self):
     @TinyJit
     def add(a, b): return None

test/backend/test_linearizer.py

@@ -1,13 +1,12 @@
 import numpy as np
 import unittest
-from dataclasses import replace
 
 from tinygrad.codegen.opt import Opt, OptOps
 from tinygrad.uop.ops import UOp, Ops, GroupOp, AxisType
 from tinygrad.device import Device, Buffer, is_dtype_supported
 from tinygrad.tensor import Tensor, _to_np_dtype
-from tinygrad.engine.realize import run_linear, CompiledRunner, get_program
-from tinygrad.schedule import linear_to_schedule
+from tinygrad.engine.realize import run_linear, CompiledRunner
+from tinygrad.codegen import to_program
 from tinygrad.helpers import Context, flatten, dedup, TC_SELECT, TC_OPT, DEV
 from tinygrad.dtype import DType, dtypes, PtrDType, AddrSpace
 from tinygrad.renderer.ptx import PTXRenderer
@@ -26,9 +25,9 @@ class TestLinearizer(unittest.TestCase):
     a, b = Tensor.randn(4).realize(), Tensor.randn(4).realize()
     np_a, np_b = a.numpy(), b.numpy()
     c = ((a.shrink(((0, 2),)) - a.shrink(((2, 4),))) - (b.shrink(((0, 2),)) - b.shrink(((2, 4),))))
-    sched = c.schedule()
-    for si in sched: si.run()
-    rawbufs = sched[-1].bufs
+    linear = c.schedule_linear()
+    run_linear(linear)
+    rawbufs = [s.buffer for s in linear.src[-1].src[1:] if s.op is not Ops.BIND]
     assert len(rawbufs) == 3 and set(rawbufs[1:]) == {a.uop.base.realized, b.uop.base.realized}
     np_c = (np_a[:2] - np_a[2:]) - (np_b[:2] - np_b[2:])
     np.testing.assert_allclose(np_c, c.numpy(), atol=1e-4, rtol=1e-4)
@@ -46,7 +45,7 @@ class TestLinearizer(unittest.TestCase):
     tst = Tensor.ones(16, dtype=dtypes.int).contiguous().realize()
     out = tst.neg().cast(dtypes.char).cast(dtypes.int).cast(dtypes.char) * 2
     ast = helper_linearizer_opt(out)
-    uops = get_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).uops
+    uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[2].src)
     self.assertEqual(len([x for x in uops if x.op is Ops.CAST]), 1)
 
   @unittest.expectedFailure
@@ -54,7 +53,7 @@ class TestLinearizer(unittest.TestCase):
     tst = Tensor.ones(16, dtype=dtypes.int).contiguous().realize()
     out = tst.neg().cast(dtypes.char).cast(dtypes.int) * 2
     ast = helper_linearizer_opt(out)
-    uops = get_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).uops
+    uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[2].src)
     self.assertEqual(len([x for x in uops if x.op is Ops.CAST]), 0)
 
   @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "broken on ptx")
@@ -64,7 +63,7 @@ class TestLinearizer(unittest.TestCase):
     b = Tensor.empty(16)
     out = img.conv2d(w, b)
     ast = helper_linearizer_opt(out)
-    uops = get_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).uops
+    uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[2].src)
     # slice at the last loop end
     uslice = [i for i,u in enumerate(uops) if u.op == Ops.END][-1]
     # only valid test if outermost range is the reduce
@@ -85,7 +84,7 @@ class TestLinearizer(unittest.TestCase):
     a = Tensor.randn(2, ).realize()
     out = a.reshape(2, 1).expand(2, 3).sum()
     ast = helper_linearizer_opt(out, wanna_output=[np.broadcast_to(a.numpy().reshape(2, 1), (2, 3)).sum()])
-    uops = get_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).uops
+    uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[2].src)
     ranges = [i for i,u in enumerate(uops) if u.op is Ops.RANGE]
     assert len(ranges) == 1 # NOTE: it collapses now
 
@@ -93,7 +92,7 @@ class TestLinearizer(unittest.TestCase):
     a = Tensor.randn(2, ).realize()
     out = a.reshape(2, 1).expand(2, 3).expand(2, 2, 3).sum()
     ast = helper_linearizer_opt(out, wanna_output=[np.broadcast_to(np.broadcast_to(a.numpy().reshape(2, 1), (2, 3)), (2, 2, 3)).sum()])
-    uops = get_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).uops
+    uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[2].src)
     ranges = [i for i,u in enumerate(uops) if u.op is Ops.RANGE]
     assert len(ranges) == 1 # NOTE: it collapses now
 
@@ -101,7 +100,7 @@ class TestLinearizer(unittest.TestCase):
     a = Tensor([2, 2]).realize()
     out = a.reshape(2, 1).pad(((1, 1), (1, 1)), value=2).sum()
     ast = helper_linearizer_opt(out, wanna_output=[24])
-    uops = get_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).uops
+    uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[2].src)
     ranges = [i for i,u in enumerate(uops) if u.op is Ops.RANGE]
     # RANGE -> ALU -> RANGE -> ALU + LOAD -> STORE
     assert any(x.op in GroupOp.ALU for x in uops[ranges[0]:ranges[1]])
@@ -114,7 +113,7 @@ class TestLinearizer(unittest.TestCase):
     b = Tensor.randn(1, 1).realize()
     out = (a + b[0]).sum() + b[0]
     ast = helper_linearizer_opt(out, wanna_output=[(a.numpy()+b.numpy()[0]).sum()+b.numpy()])
-    uops = get_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).uops
+    uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[2].src)
     ranges = [i for i,u in enumerate(uops) if u.op is Ops.RANGE]
     # LOAD -> RANGE -> LOAD -> STORE
     assert len([x for x in uops[:ranges[0]] if x.op is Ops.LOAD]) == 1
@@ -124,7 +123,7 @@ class TestLinearizer(unittest.TestCase):
     b = Tensor.randn(1, 1).realize()
     out = (a.reshape(2, 1).expand(2, 3) + b[0]).sum() + b[0]
     ast = helper_linearizer_opt(out, wanna_output=[(np.broadcast_to(a.numpy().reshape(2, 1), (2, 3)) + b.numpy()[0]).sum() + b.numpy()])
-    uops = get_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).uops
+    uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[2].src)
     ranges = [i for i,u in enumerate(uops) if u.op is Ops.RANGE]
     assert len(ranges) == 1 # NOTE: it collapses now
 
@@ -135,7 +134,8 @@ class TestLinearizer(unittest.TestCase):
     # these are of size 3 to avoid float4 coalesce
     r = a[:-1] + a[1:]
 
-    uops = get_program(replace_opts(r.schedule()[-1].ast, [Opt(op=OptOps.UPCAST, axis=0, arg=0)]), renderer=Device[Device.DEFAULT].renderer).uops
+    uops = tuple(to_program(replace_opts(r.schedule_linear().src[-1].src[0], [Opt(op=OptOps.UPCAST, axis=0, arg=0)]),
+                       renderer=Device[Device.DEFAULT].renderer).src[2].src)
     num_loads = len([uop for uop in uops if uop.op is Ops.LOAD])
     assert num_loads <= 4, "more load uops than needed"
     assert num_loads >= 4, "unexpected number of uops, maybe this test needs updating?"
@@ -147,7 +147,8 @@ class TestLinearizer(unittest.TestCase):
     a, b = Tensor.randn(1).realize(), Tensor.randn(1).realize()
     r = a.expand([2]) + b.expand([2])
 
-    uops = get_program(replace_opts(r.schedule()[-1].ast, [Opt(op=OptOps.UPCAST, axis=0, arg=0)]), renderer=Device[Device.DEFAULT].renderer).uops
+    uops = tuple(to_program(replace_opts(r.schedule_linear().src[-1].src[0], [Opt(op=OptOps.UPCAST, axis=0, arg=0)]),
+                       renderer=Device[Device.DEFAULT].renderer).src[2].src)
     num_ops = len([uop for uop in uops if uop.op in GroupOp.ALU])
     assert num_ops <= 1, "more alu uops than needed"
 
@@ -156,8 +157,8 @@ class TestLinearizer(unittest.TestCase):
     x, w = Tensor.randn((1,1,3)).realize(), Tensor.randn((1,1,2)).realize()
     r = Tensor.conv2d(x,w,padding=1).relu()
 
-    uops = get_program(replace_opts(r.schedule()[-1].ast, [Opt(op=OptOps.UPCAST, axis=0, arg=0), Opt(op=OptOps.UNROLL, axis=0, arg=0)]),
-                       renderer=Device[Device.DEFAULT].renderer).uops
+    uops = tuple(to_program(replace_opts(r.schedule_linear().src[-1].src[0],
+      [Opt(op=OptOps.UPCAST, axis=0, arg=0), Opt(op=OptOps.UNROLL, axis=0, arg=0)]), renderer=Device[Device.DEFAULT].renderer).src[2].src)
     accs = [u for u in uops if u.op is Ops.DEFINE_REG]
     stores = [u for u in uops if u.op is Ops.STORE]
     assert len(accs) == 0  # it's removed now
@@ -169,8 +170,9 @@ class TestLinearizer(unittest.TestCase):
   @unittest.skipUnless(Device.DEFAULT == "CPU", "test only for CPU")
   def test_upcast_with_locals_cpu(self):
     out = Tensor.ones(64,64).contiguous() @ Tensor.ones(64,64).contiguous()
-    prg = get_program(replace_opts(out.schedule()[-1].ast, [Opt(OptOps.LOCAL, axis=0, arg=4)]), renderer=Device[Device.DEFAULT].renderer).uops
-    self.assertEqual(len(prg.src.split("for")), 5)
+    prg = to_program(replace_opts(out.schedule_linear().src[-1].src[0], [Opt(OptOps.LOCAL, axis=0, arg=4)]),
+                      renderer=Device[Device.DEFAULT].renderer)
+    self.assertEqual(len(prg.src[3].arg.split("for")), 5)
 
   @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals")
   @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_shared, "test requires shared")
@@ -180,9 +182,9 @@ class TestLinearizer(unittest.TestCase):
     x, y = Tensor.rand(1,128), Tensor.rand(128, 128)
     r = (x@y).relu()
     opts_to_apply = [Opt(op=OptOps.GROUP, axis=0, arg=8), Opt(op=OptOps.LOCAL, axis=0, arg=4), Opt(op=OptOps.UPCAST, axis=0, arg=4)]
-    program = get_program(replace_opts(r.schedule()[-1].ast, opts_to_apply), renderer=Device[Device.DEFAULT].renderer)
+    program = to_program(replace_opts(r.schedule_linear().src[-1].src[0], opts_to_apply), renderer=Device[Device.DEFAULT].renderer)
 
-    stores = [u for u in program.uops if u.op is Ops.STORE and u.src[0].dtype.addrspace != AddrSpace.REG]
+    stores = [u for u in tuple(program.src[2].src) if u.op is Ops.STORE and u.src[0].dtype.addrspace != AddrSpace.REG]
 
     # the first store is to lds and can be upcasted
     assert stores[0].src[1].dtype == dtypes.float.vec(4)
@@ -194,7 +196,8 @@ class TestLinearizer(unittest.TestCase):
   def test_zero_fold(self):
     a, b = Tensor.randn(1).realize(), Tensor.randn(1).realize()
     r = Tensor.stack(a, b)
-    uops = get_program(replace_opts(r.schedule()[-1].ast, [Opt(op=OptOps.UPCAST, axis=0, arg=0)]), renderer=Device[Device.DEFAULT].renderer).uops
+    uops = tuple(to_program(replace_opts(r.schedule_linear().src[-1].src[0], [Opt(op=OptOps.UPCAST, axis=0, arg=0)]),
+                       renderer=Device[Device.DEFAULT].renderer).src[2].src)
     num_ops = len([uop for uop in uops if uop.op in GroupOp.ALU])
     assert num_ops == 0, "more alu uops than needed"
 
@@ -203,16 +206,16 @@ class TestLinearizer(unittest.TestCase):
       (dtypes.bool, dtypes.int), (dtypes.int16, dtypes.int), (dtypes.float16, dtypes.float), (dtypes.bfloat16, dtypes.float)):
       if is_dtype_supported(tensor_dtype) and is_dtype_supported(acc_dtype):
         a = Tensor([1, 2, 3], dtype=tensor_dtype).sum()
-        realized_ast = a.schedule()[-1].ast
-        program = get_program(replace_opts(realized_ast, []), renderer=Device[Device.DEFAULT].renderer)
-        local = [uop for uop in program.uops if uop.op is Ops.DEFINE_REG]
+        realized_ast = a.schedule_linear().src[-1].src[0]
+        program = to_program(replace_opts(realized_ast, []), renderer=Device[Device.DEFAULT].renderer)
+        local = [uop for uop in tuple(program.src[2].src) if uop.op is Ops.DEFINE_REG]
         assert local[0].dtype.base == acc_dtype
 
   def test_arg_acc_dtype(self):
     def helper_arg_acc_dtype(c: Tensor, expected_dtype:DType):
-      realized_ast = c.schedule()[-1].ast
-      program = get_program(replace_opts(realized_ast, []), renderer=Device[Device.DEFAULT].renderer)
-      local = [uop for uop in program.uops if uop.op is Ops.DEFINE_REG]
+      realized_ast = c.schedule_linear().src[-1].src[0]
+      program = to_program(replace_opts(realized_ast, []), renderer=Device[Device.DEFAULT].renderer)
+      local = [uop for uop in tuple(program.src[2].src) if uop.op is Ops.DEFINE_REG]
       self.assertEqual(local[0].dtype.base, expected_dtype)
 
     tests = (
@@ -239,7 +242,7 @@ class TestLinearizer(unittest.TestCase):
     opt = [Opt(OptOps.UNROLL, 0, 4), Opt(OptOps.UPCAST, 0, 4)]
     ast = helper_linearizer_opt(r, [opt])
     # the uops graph is DEFINE_REG -> 4x STORE 0.0 -> RANGE -> 4x ALU -> 4x STORE -> ENDRANGE
-    uops = get_program(replace_opts(ast, opt), renderer=Device[Device.DEFAULT].renderer).uops
+    uops = tuple(to_program(replace_opts(ast, opt), renderer=Device[Device.DEFAULT].renderer).src[2].src)
     begin_range = [i for i, x in enumerate(uops) if x.op is Ops.RANGE][-1]
     end_range = [i for i, x in enumerate(uops) if x.op is Ops.END][0]
     for i,u in enumerate(uops): print(i, u.op, [uops.index(s) for s in u.src], u.arg, u.dtype)
@@ -259,7 +262,7 @@ class TestLinearizer(unittest.TestCase):
     # shrink so that the dims do not collapse
     t = Tensor.ones(5, 6, 7).contiguous().realize().shrink(((0, 4), (0, 5), (0, 6)))
     ast = helper_linearizer_opt(t+1)
-    uops = get_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).uops
+    uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[2].src)
     idxs = dedup([uop for uop in uops if uop.op is Ops.SPECIAL])
     idxs = sorted(idxs, key=lambda uop: uop.arg)
     assert (idxs[0].arg, idxs[0].src[0].arg) == ('gidx0', 6), idxs[0]
@@ -268,10 +271,10 @@ class TestLinearizer(unittest.TestCase):
 
   def test_sum_collapse(self):
     t = Tensor([2]).reshape(1, 1).expand(256, 256).sum()
-    sched = [si for si in t.schedule() if si.ast.op is Ops.SINK]
+    sched = [si for si in t.schedule_linear().src if si.src[0].op is Ops.SINK]
     # sum_collapse is a full collapse now
     assert len(sched) == 1
-    assert not any(u.op is Ops.REDUCE_AXIS for u in sched[0].ast.toposort()), "found reduce in sum collapse"
+    assert not any(u.op is Ops.REDUCE_AXIS for u in sched[0].src[0].toposort()), "found reduce in sum collapse"
     #lin = Kernel(sched[0].ast)
     #assert not any(u.op is Ops.RANGE for u in lin.linearize().uops), "found loop in sum collapse"
 
@@ -288,17 +291,16 @@ class TestLinearizer(unittest.TestCase):
     b = a.shrink(((1, 2), None)).pad(((1, 2), None))
     a.assign(b.where(2, a))
     linear, var_vals = a.linear_with_vars()
-    sched_copy = linear_to_schedule(linear)
-    assert len(sched_copy) == 1
+    assert len(linear.src) == 1
     run_linear(linear, var_vals)
     np.testing.assert_equal(a.flatten().numpy(), [1.,1.,1.,1.,2.,2.,2.,2.,1.,1.,1.,1.,1.,1.,1.,1.])
-    program = get_program(replace_opts(sched_copy[-1].ast, []), renderer=Device[Device.DEFAULT].renderer)
-    assert not any(u.op == Ops.WHERE for u in program.uops), "found where where where should be folded"
+    program = to_program(replace_opts(linear.src[-1].src[0], []), renderer=Device[Device.DEFAULT].renderer)
+    assert not any(u.op == Ops.WHERE for u in tuple(program.src[2].src)), "found where where where should be folded"
 
   def test_phi_simplification(self):
     def helper(t, max_ops=0):
       ast = helper_linearizer_opt(t)
-      uops = get_program(ast, renderer=Device[Device.DEFAULT].renderer).uops
+      uops = tuple(to_program(ast, renderer=Device[Device.DEFAULT].renderer).src[2].src)
       # ignore kernel optimized IF statements for now
       if if_op:=next((u for u in uops if u.op is Ops.IF), None):
         uops = uops[:uops.index(if_op)]
@@ -330,7 +332,7 @@ class TestLinearizer(unittest.TestCase):
     out = x.matmul(y)
     with Context(TC=0):
       ast = helper_linearizer_opt(out)
-      uops = get_program(ast, renderer=Device[Device.DEFAULT].renderer).uops
+      uops = tuple(to_program(ast, renderer=Device[Device.DEFAULT].renderer).src[2].src)
     # check that the float4 cast collapses
     store_vals = [u.src[1] for u in uops if u.op is Ops.STORE and u.src[0].dtype.addrspace != AddrSpace.REG]
     for val in store_vals:
@@ -341,7 +343,7 @@ class TestLinearizer(unittest.TestCase):
     x = Tensor.randn((4,3,6,6)).realize()
     out = x.flip((0,1)).contiguous()
     ast = helper_linearizer_opt(out)
-    store_val = [u.src[1] for u in get_program(ast, renderer=Device[Device.DEFAULT].renderer).uops if u.op is Ops.STORE][0]
+    store_val = [u.src[1] for u in tuple(to_program(ast, renderer=Device[Device.DEFAULT].renderer).src[2].src) if u.op is Ops.STORE][0]
     assert store_val.dtype == dtypes.float.vec(4) and store_val.op is not Ops.STACK
 
   @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals")
@@ -354,7 +356,7 @@ class TestLinearizer(unittest.TestCase):
             Opt(OptOps.UNROLL, 0, 4), Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UPCAST, 1, 2)] # upcast accs in both reduces
     ast = helper_linearizer_opt(out, opts=[opt])
     def get_recursive(uop): return set.union(set(uop.src), [uop], *[get_recursive(v) for v in uop.src])
-    uops = get_program(replace_opts(ast, opt), renderer=Device[Device.DEFAULT].renderer).uops
+    uops = tuple(to_program(replace_opts(ast, opt), renderer=Device[Device.DEFAULT].renderer).src[2].src)
     local_stores = [u for u in uops if u.op is Ops.STORE and any(x.op is Ops.DEFINE_LOCAL for x in get_recursive(u.src[0]))]
     global_stores = [u for u in uops if u.op is Ops.STORE and any(x.op is Ops.PARAM for x in get_recursive(u.src[0]))]
     barrier = [u for u in uops if u.op is Ops.BARRIER]
@@ -375,7 +377,7 @@ class TestLinearizer(unittest.TestCase):
     x, y = Tensor.rand(1,128), Tensor.rand(128, 128)
     r = (x@y).relu()
     ast = helper_linearizer_opt(r)
-    uops = get_program(ast, renderer=Device[Device.DEFAULT].renderer).uops
+    uops = tuple(to_program(ast, renderer=Device[Device.DEFAULT].renderer).src[2].src)
     stores = [u for u in uops if u.op is Ops.STORE and u.src[0].dtype.addrspace != AddrSpace.REG]
 
     # the float4 value stores directly in lds and we skip upcast
@@ -390,15 +392,17 @@ class TestLinearizer(unittest.TestCase):
 def helper_realized_ast(r:Tensor|list[Tensor]) -> tuple[UOp, list[Buffer]]:
   if isinstance(r, Tensor): r = [r]
   linear, var_vals = Tensor.linear_with_vars(*r)
-  s = linear_to_schedule(linear)
   run_linear(UOp(Ops.LINEAR, src=linear.src[:-1]), var_vals)  # run all kernels except the last one
-  assert s[-1].ast.op is Ops.SINK, f"helper_realized_ast expects a SINK {s[-1]}"
-  # now all input buffers in s[-1] should be realized
+  last_call = linear.src[-1]
+  ast = last_call.src[0]
+  assert ast.op is Ops.SINK, f"helper_realized_ast expects a SINK {last_call}"
+  last_bufs = [s.buffer for s in last_call.src[1:] if s.op is not Ops.BIND]
+  # now all input buffers in last_call should be realized
   # create fresh buffers for the outputs
-  bufs = [Buffer(x.device, x.size, x.dtype).allocate() if i < len(s[-1].ast.src) else x for i,x in enumerate(s[-1].bufs)]
+  bufs = [Buffer(x.device, x.size, x.dtype).allocate() if i < len(ast.src) else x for i,x in enumerate(last_bufs)]
   # ensure buffers are allocated
   for b in bufs: b.ensure_allocated()
-  return s[-1].ast, bufs
+  return ast, bufs
 
 def helper_linearizer_ast(ast:UOp, inputs:list[Tensor], *args, **kwargs):
   assert isinstance(ast, UOp), "ast must be UOp"
@@ -425,7 +429,7 @@ def _helper_linearizer_opt_ast(realized_ast:UOp, real_bufs:list[Buffer], opts=[]
 
   def get_prg(opts):
     ast = realized_ast if opts is None else replace_opts(realized_ast, list(opts))
-    return CompiledRunner(replace(get_program(ast, renderer=Device[Device.DEFAULT].renderer), device=device))
+    return CompiledRunner(to_program(ast, renderer=Device[Device.DEFAULT].renderer), device)
 
   def check_opt(opts):
     prg = get_prg(opts=opts)

test/backend/test_linearizer_dumb.py

@@ -6,7 +6,7 @@ import unittest
 from tinygrad import Device, dtypes
 from tinygrad.uop.ops import UOp, Ops, AxisType, KernelInfo
 from tinygrad.codegen.opt.search import Opt, OptOps
-from tinygrad.engine.realize import get_program
+from tinygrad.codegen import to_program
 
 class TestLinearizerFailure(unittest.TestCase):
   @unittest.skipUnless(Device.DEFAULT == "METAL", "only tested on METAL")
@@ -25,7 +25,7 @@ class TestLinearizerFailure(unittest.TestCase):
     c11 = c5.alu(Ops.CMPNE, ((((c3*UOp.const(dtypes.weakint, 6000))+c6)+((c7*UOp.const(dtypes.weakint, 16))+c8)).alu(Ops.CMPLT, UOp.const(dtypes.weakint, 59999)).where(UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 1)).reduce(c7, c8, arg=Ops.ADD)+UOp.const(dtypes.int, -1))).where(UOp.const(dtypes.uchar, 0), c10).reduce(c6, arg=Ops.ADD)
     c12 = c0.index((((c1*UOp.const(dtypes.weakint, 7840))+(c2*UOp.const(dtypes.weakint, 10)))+c3).valid(UOp.const(dtypes.bool, True))).store(c11).end(c1, c2, c3)
     ast = c12.sink(arg=KernelInfo(name='test', axis_types=(), dont_use_locals=False, applied_opts=(Opt(op=OptOps.GROUP, axis=1, arg=16),), opts_to_apply=None))
-    _ = get_program(ast, Device["METAL"].renderer)
+    _ = to_program(ast, Device["METAL"].renderer)
 
 if __name__ == '__main__':
   unittest.main()

test/backend/test_multitensor.py

@@ -4,8 +4,7 @@ from tinygrad.device import is_dtype_supported
 from tinygrad.uop.ops import Ops, UOp
 from tinygrad.helpers import getenv, prod, Context
 from tinygrad.nn.state import get_parameters, get_state_dict
-from tinygrad.engine.realize import CompiledRunner, run_linear
-from tinygrad.schedule import linear_to_schedule
+from tinygrad.engine.realize import run_linear, compile_linear
 import numpy as np
 from hypothesis import given, strategies as strat, settings
 from test.helpers import not_support_multi_device, needs_second_gpu, slow, call_is_graph
@@ -128,12 +127,9 @@ class TestMultiTensor(unittest.TestCase):
     X = Tensor.ones(256).contiguous().realize()
     X.shard_(devices_2, 0)
     out = (X + X)
-    sched = out.schedule()
-    names = []
-    for si in sched:
-      si.lower()
-      if isinstance(si.prg, CompiledRunner): names.append(si.prg.p.name)
-      si.run()
+    linear = compile_linear(out.schedule_linear())
+    names = [call.src[0].src[0].arg.name for call in linear.src if call.src[0].op is Ops.PROGRAM]
+    run_linear(linear)
     self.assertEqual(len(set(names)), 1, "function was relinearized")
 
   def test_shard_same_device(self):
@@ -194,9 +190,9 @@ class TestMultiTensor(unittest.TestCase):
     for i in range(2):
       xt = X[i*2:i*2+2].contiguous()
       linear, var_vals = xt.linear_with_vars()
-      #kernels = [s for s in linear_to_schedule(linear) if s.ast.op is Ops.SINK]
+      #kernels = [call for call in linear.src if call.src[0].op is Ops.SINK]
       #self.assertEqual(len(kernels), 1)
-      #self.assertEqual(kernels[0].bufs[0].device, devices_2[i])
+      #self.assertEqual(kernels[0].src[1].buffer.device, devices_2[i])
       run_linear(linear, var_vals)
       np.testing.assert_equal(xt.numpy(), X_np[i*2:i*2+2])
 
@@ -555,6 +551,21 @@ class TestMultiTensor(unittest.TestCase):
       np.testing.assert_allclose(c.numpy(), a.numpy()+b.numpy(), atol=1e-4, rtol=1e-5)
     assert jf.captured is not None
 
+  def test_multi_tensor_jit_graph_assign_updates_each_shard(self):
+    @TinyJit
+    def jf(out: Tensor) -> Tensor:
+      tmp = (Tensor.arange(4, dtype=dtypes.float).shard(devices_2, 0) + 1).contiguous().realize()
+      out.assign((tmp + 1).contiguous()).realize()
+      return out
+
+    out = Tensor.full((4,), -1.0).shard(devices_2, 0).contiguous().realize()
+    expected = np.arange(4, dtype=np.float32) + 2
+    for _ in range(5):
+      out.assign(Tensor.full((4,), -1.0).shard(devices_2, 0).contiguous()).realize()
+      jf(out)
+      np.testing.assert_allclose(out.numpy(), expected, atol=1e-4, rtol=1e-5)
+    assert jf.captured is not None
+
   def test_multi_tensor_jit_body(self):
     @TinyJit
     def jf() -> Tensor:
@@ -640,18 +651,15 @@ class TestMultiTensor(unittest.TestCase):
     for p in get_parameters(bn): p.shard_(devices_4).realize()
 
     out = bn(t)
-    scheds = [sched for sched in out.schedule() if sched.bufs[0].device in devices_4 and sched.ast.op is not Ops.COPY]
-    assert set(sched.bufs[0].device for sched in scheds) == set(devices_4), "should have ast on each shard device"
-    asts = [sched.ast for sched in scheds]
-    self.assertEqual(len(asts), 4)
-    # ast are the same on devices
-    self.assertEqual(len(set(asts)), 1)
+    scheds = [call for call in out.schedule_linear().src if call.src[0].op is not Ops.COPY and set(call.device) <= set(devices_4)]
+    self.assertEqual(set(scheds[0].device), set(devices_4), "should have ast on each shard device")
+    self.assertEqual(len(set(s.src[0] for s in scheds)), 1)
 
   def test_flip(self):
     rng = Tensor.rand((10, 10, 10))
     t0 = rng.shard(devices_2, axis=1)
     out = t0.flip(0) + 1
-    self.assertTrue((rng.flip(0)+1).allclose(out.to(rng.device)))
+    self.assertTrue((rng.flip(0)+1).allclose(out.to(rng.device)).item())
 
   @unittest.skip("flaky")
   def test_reshape_on_axis(self):
@@ -685,7 +693,7 @@ class TestMultiTensor(unittest.TestCase):
 
     # test no left join
     with self.assertRaises((AssertionError, ValueError)):
-      t0.reshape((26*15,7)).contiguous().schedule()
+      t0.reshape((26*15,7)).contiguous().schedule_linear()
 
   # it doesn't work like this anymore
   # NOTE: this never failed in assign_multi, it failed tensor spec because MULTI was never pushed in the graph
@@ -696,7 +704,7 @@ class TestMultiTensor(unittest.TestCase):
     with self.assertRaises(RuntimeError):
       # don't allow assigns that change axes
       t_none.assign(t_zero)
-      t_none.schedule()
+      t_none.schedule_linear()
 
   def test_init_rand_with_multiple_devices_fail(self):
     # init rand with multi device is not allowed
@@ -794,7 +802,7 @@ class TestMultiTensor(unittest.TestCase):
     t = Tensor.ones(16, 16, dtype=dtypes.int).shard(devices_2, axis=0)
     out = Tensor.full_like(t, 2)[:, :8]
     linear, var_vals = out.linear_with_vars()
-    self.assertEqual(len(linear_to_schedule(linear)), 0)
+    self.assertEqual(len(linear.src), 0)
     run_linear(linear, var_vals)
     self.assertEqual(out.tolist(), [[2]*8]*16)
 
@@ -837,7 +845,7 @@ class TestMultiTensor(unittest.TestCase):
       a = Tensor.arange(3).realize()
       zeros = Tensor.zeros(3).realize()
     b = a.to(devices_2)*zeros.to(devices_2)
-    sched = b.schedule()
+    sched = b.schedule_linear().src
     self.assertEqual(len(sched), 0)
     self.assertListEqual(b.tolist(), [0, 0, 0])
 
@@ -848,7 +856,7 @@ class TestHandleData(unittest.TestCase):
     device = (d0, d1, d2, d3)
     t = Tensor([1, 2, 3, 4]).shard(device).realize()
     not_covered = t.to(d5)
-    sched = not_covered.schedule()
+    sched = not_covered.schedule_linear().src
     assert len(sched) == 1
     # setup again because create_schedule has side effect
     t = Tensor([1, 2, 3, 4]).shard(device).realize()
@@ -858,7 +866,7 @@ class TestHandleData(unittest.TestCase):
     for d in device:
       t = Tensor([1, 2, 3, 4]).shard(device).realize()
       covered = t.to(d)
-      sched = covered.schedule()
+      sched = covered.schedule_linear().src
       # TODO: this isn't optimized out anymore
       #assert len(sched) == 0
       # setup again because create_schedule has side effect
@@ -879,18 +887,18 @@ class TestShrinkMultiTensorShardedAxis(unittest.TestCase):
     with self.assertRaises(AssertionError):
       # sharded axis shrink on non-device boundry is not allowed
       a = t.shrink(((0, 3), (0, 8))).contiguous()
-      a.schedule()
+      a.schedule_linear()
     a = t.shrink(((0, 2), (2, 4)))
     assert a.shape == (2, 2)
     ref = Tensor.arange(64).reshape(8, 8).shrink(((0, 2), (2, 4)))
     np.testing.assert_equal(a.numpy(), ref.numpy())
 
     a = t.shrink(((0, 2), (0, 8))).contiguous()
-    a.schedule()
+    a.schedule_linear()
     assert a.shape == (2, 8)
 
     p = a.pad(((0, 6), (0, 0))).contiguous()
-    p.schedule()
+    p.schedule_linear()
     assert p.shape == (8, 8)
 
   @given(strat.sampled_from([dtypes.float, dtypes.int, dtypes.int64, dtypes.int16]))
@@ -1105,9 +1113,9 @@ class TestBatchNorm(unittest.TestCase):
           p.to_(devices)
 
       synced_out = synced_bn(x)
-      synced_si = list(synced_out.schedule())
+      synced_si = list(synced_out.schedule_linear().src)
       unsynced_out = unsynced_bn(x)
-      unsynced_si = list(unsynced_out.schedule())
+      unsynced_si = list(unsynced_out.schedule_linear().src)
 
     # TODO: test synced / unsynced batchnorm cross device kernel and copies
     assert synced_si
@@ -1144,12 +1152,12 @@ class TestMultiBufferView(unittest.TestCase):
   def setUp(self): pass
 
   def _check(self, a_ref:Tensor, a_multi:Tensor, view_fn):
-    """Apply view_fn to both, verify zero compiled kernels and matching values."""
     b_ref = view_fn(a_ref)
     b_multi = view_fn(a_multi).contiguous()
     linear, var_vals = b_multi.linear_with_vars()
-    compiled = [si for si in linear_to_schedule(linear) if isinstance(si.prg, CompiledRunner)]
-    self.assertEqual(len(compiled), 0, f"expected zero compiled kernels, got {len(compiled)}")
+    if all(hasattr(Device[d].allocator, "_offset") for d in b_multi.device):
+      compiled = [call for call in linear.src if call.src[0].op is Ops.SINK]
+      self.assertEqual(len(compiled), 0, f"expected zero compiled kernels, got {len(compiled)}")
     run_linear(linear, var_vals)
     np.testing.assert_equal(b_multi.numpy(), b_ref.numpy())
 
@@ -1177,11 +1185,13 @@ class TestMultiBufferView(unittest.TestCase):
   def test_4_devices(self):
     ref = Tensor.arange(8*12).reshape(8, 12).contiguous().realize()
     a = Tensor.arange(8*12).reshape(8, 12).contiguous().shard(devices_4, axis=1).realize()
-    linear, var_vals = a[5].contiguous().linear_with_vars()
-    compiled = [si for si in linear_to_schedule(linear) if isinstance(si.prg, CompiledRunner)]
-    self.assertEqual(len(compiled), 0)
+    out = a[5].contiguous()
+    linear, var_vals = out.linear_with_vars()
+    if all(hasattr(Device[d].allocator, "_offset") for d in out.device):
+      compiled = [call for call in linear.src if call.src[0].op is Ops.SINK]
+      self.assertEqual(len(compiled), 0)
     run_linear(linear, var_vals)
-    np.testing.assert_equal(a[5].contiguous().numpy(), ref[5].numpy())
+    np.testing.assert_equal(out.numpy(), ref[5].numpy())
 
 @unittest.skipIf(not_support_multi_device(), "need multi")
 class TestMultiFromUnrenderable(unittest.TestCase):

test/backend/test_nn.py

@@ -9,7 +9,6 @@ from tinygrad.nn import Conv1d, ConvTranspose1d, Conv2d, ConvTranspose2d, Linear
 from tinygrad.nn import BatchNorm, LayerNorm, LayerNorm2d, GroupNorm, InstanceNorm, RMSNorm, LSTMCell
 from tinygrad.nn.state import load_state_dict
 from tinygrad.engine.realize import run_linear
-from tinygrad.schedule import linear_to_schedule
 from test.helpers import not_support_multi_device, needs_second_gpu, slow
 
 @slow
@@ -433,7 +432,7 @@ class TestNN(unittest.TestCase):
                 [12, 19, 8, 1]])
     result = layer(a)
     linear, var_vals = result.linear_with_vars()
-    self.assertEqual(len([item for item in linear_to_schedule(linear) if item.ast.op is Ops.SINK]), kcount,
+    self.assertEqual(len([call for call in linear.src if call.src[0].op is Ops.SINK]), kcount,
                      "first run realizes weight and embedding")
     run_linear(linear, var_vals)
 
@@ -442,7 +441,7 @@ class TestNN(unittest.TestCase):
                 [7, 8, 9]])
     result = layer(b)
     linear, var_vals = result.linear_with_vars()
-    self.assertEqual(1, len([item for item in linear_to_schedule(linear) if item.ast.op is Ops.SINK]),
+    self.assertEqual(1, len([call for call in linear.src if call.src[0].op is Ops.SINK]),
                      "second run realizes embedding only")
     run_linear(linear, var_vals)
     print(f"Embedding used {GlobalCounters.global_ops} ops")

test/backend/test_opt_gemm.py

@@ -1,10 +1,10 @@
 import numpy as np
 import unittest
-from tinygrad import Tensor, Device
+from tinygrad import Tensor
 from tinygrad.helpers import get_single_element
 from tinygrad.codegen.opt import Opt, OptOps
-from tinygrad.engine.realize import CompiledRunner, get_program
-from tinygrad.schedule import ExecItem
+from tinygrad.engine.realize import run_linear
+from tinygrad.uop.ops import Ops, UOp
 from test.helpers import replace_opts
 
 class TestOptGemm(unittest.TestCase):
@@ -19,10 +19,10 @@ class TestOptGemm(unittest.TestCase):
   def _test_gemm_unrolled_permute_l(self, opts=[]):
     t = self.a.T @ self.b.T
     # TODO: this should be a generic test helper
-    si = get_single_element(t.schedule())
-    run = CompiledRunner(get_program(replace_opts(si.ast, opts), renderer=Device[Device.DEFAULT].renderer))
-    ExecItem(si.ast, list(si.bufs), prg=run).run()
-    test = si.bufs[0].numpy().reshape(self.res.shape)
+    call = get_single_element(t.schedule_linear().src)
+    new_call = call.replace(src=(replace_opts(call.src[0], opts), *call.src[1:]))
+    run_linear(UOp(Ops.LINEAR, src=(new_call,)))
+    test = call.src[1].buffer.numpy().reshape(self.res.shape)
     np.testing.assert_allclose(self.res, test, atol=1e-4)
 
   def test_gemm_unrolled_permute_l_44(self):

test/backend/test_pickle.py

@@ -142,10 +142,10 @@ class TestPickle(unittest.TestCase):
   def test_pickle_schedule(self):
     a = Tensor([1,2])
     out = a + 2
-    sched = out.schedule()
+    sched = out.schedule_linear()
     pk = pickle.dumps(sched)
     sched_pk = pickle.loads(pk)
-    self.assertEqual(sched_pk[-1].ast, sched[-1].ast)
+    self.assertEqual(sched_pk.src[-1].src[0], sched.src[-1].src[0])
 
   def test_pickle_renderer(self):
     from tinygrad.device import Device

test/backend/test_profiler.py

@@ -44,9 +44,9 @@ class TestProfiler(unittest.TestCase):
 
     TestProfiler.a = Tensor([0.,1.], device=Device.DEFAULT).realize()
     TestProfiler.b = self.a + 1
-    si = self.b.schedule()[-1]
+    si = self.b.schedule_linear().src[-1]
 
-    TestProfiler.runner = get_runner(TestProfiler.d0.device, si.ast)
+    TestProfiler.runner = get_runner(TestProfiler.d0.device, si.src[0])
     TestProfiler.b.uop.buffer.allocate()
 
   def test_profile_kernel_run(self):

test/backend/test_quantize_onnx.py

@@ -1,12 +1,11 @@
 # ruff: noqa: E501
 import numpy as np
 import unittest
-from dataclasses import replace
-from tinygrad import Tensor, Context, Device, dtypes
+from tinygrad import Tensor, Context, Device, dtypes, UOp
 from tinygrad.uop.ops import Ops
 from tinygrad.codegen.opt import Opt, OptOps
-from tinygrad.engine.realize import CompiledRunner, get_program
-from tinygrad.schedule import ExecItem
+from tinygrad.engine.realize import run_linear
+from tinygrad.codegen import to_program
 from test.helpers import replace_opts
 
 N = 512
@@ -39,13 +38,16 @@ def create_gemm_model(model_path:str, batch_size=N, in_size=N, out_size=N, bias=
   return model_path
 
 def sexec(out:Tensor, opts:list[Opt], replace_src=None, run_count=3):
-  si = out.schedule()[-1]
-  prg = get_program(replace_opts(si.ast, opts), renderer=Device[Device.DEFAULT].renderer)
+  linear = out.schedule_linear()
+  call = linear.src[-1]
+  prg = to_program(replace_opts(call.src[0], opts), renderer=Device[Device.DEFAULT].renderer)
   if replace_src is not None:
-    old_name = prg.src.split("__attribute__((noinline)) void ")[1].split("(")[0]
-    prg = replace(prg, src=replace_src + "/* DSP boilerplate */" + prg.src.split("/* DSP boilerplate */")[1].replace(old_name, "fxn"))
-  new_si = ExecItem(si.ast, [x.ensure_allocated() for x in si.bufs], si.metadata, prg=CompiledRunner(prg))
-  for _ in range(run_count): new_si.run(wait=True)
+    old_name = prg.src[3].arg.split("__attribute__((noinline)) void ")[1].split("(")[0]
+    new_src = replace_src + "/* DSP boilerplate */" + prg.src[3].arg.split("/* DSP boilerplate */")[1].replace(old_name, "fxn")
+    # drop BINARY and replace SOURCE so run_linear recompiles
+    prg = prg.replace(src=prg.src[:3] + (UOp(Ops.SOURCE, arg=new_src),))
+  linear = linear.replace(src=linear.src[:-1] + (call.replace(src=(prg, *call.src[1:])),))
+  for _ in range(run_count): run_linear(linear)
 
 def get_quantized_model(sz):
   from onnxruntime.quantization import quantize_static, QuantFormat, QuantType, CalibrationDataReader
@@ -75,9 +77,9 @@ class TestQuantizeOnnxCPU(unittest.TestCase):
     run_onnx = OnnxRunner(out_file)
     inp = Tensor(np.random.uniform(size=(sz, sz)).astype(np.float32))
     with Context(QUANTIZE=1):
-      sched = run_onnx({"input":inp})["output"].schedule()
-      sched[-2].lower()
-      daccs = [u for u in sched[-2].prg.p.uops if u.op is Ops.DEFINE_REG]
+      linear = run_onnx({"input":inp})["output"].schedule_linear()
+      prg = to_program(linear.src[-2].src[0], renderer=Device[Device.DEFAULT].renderer)
+      daccs = [u for u in tuple(prg.src[2].src) if u.op is Ops.DEFINE_REG]
       assert all(u.dtype.scalar() is dtypes.int for u in daccs)
 
 @unittest.skipIf(Device.DEFAULT != "DSP", "only tests for DSP")

test/backend/test_randomness.py

@@ -4,7 +4,9 @@ from functools import partial
 from tinygrad import nn, dtypes, Tensor, Device, TinyJit, Variable
 from tinygrad.helpers import getenv, CI, OSX
 from tinygrad.device import is_dtype_supported
-from tinygrad.engine.realize import CompiledRunner
+from tinygrad.codegen import to_program
+
+from tinygrad.uop.ops import Ops
 from tinygrad.renderer.ptx import PTXRenderer
 from tinygrad.renderer.nir import NIRRenderer
 from test.helpers import not_support_multi_device, needs_second_gpu
@@ -117,12 +119,13 @@ class TestRandomness(unittest.TestCase):
 
   @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, (NIRRenderer, PTXRenderer)), "PTX and NIR use pointer arithmetic")
   def test_threefry_doesnt_use_long(self):
-    sched = Tensor.rand(20).schedule()
-    for si in sched:
-      si.lower()
-      if isinstance(si.prg, CompiledRunner):
-        for u in si.prg.p.uops:
-          self.assertNotIn(u.dtype, {dtypes.long, dtypes.ulong}, msg=f"long found in {si.prg.p.name}")
+    linear = Tensor.rand(20).schedule_linear()
+    for call in linear.src:
+      ast = call.src[0]
+      if ast.op is Ops.SINK:
+        prg = to_program(ast, renderer=Device[Device.DEFAULT].renderer)
+        for u in tuple(prg.src[2].src):
+          self.assertNotIn(u.dtype, {dtypes.long, dtypes.ulong}, msg=f"long found in {prg.arg.name}")
 
   def test_threefry_against_reference_full(self):
     Tensor.manual_seed(1337)
@@ -187,24 +190,24 @@ class TestRandomness(unittest.TestCase):
 
     Tensor.rand(1).realize()
 
-    s = Tensor.rand(20).schedule()
-    s2 = Tensor.rand(20).schedule()
+    s = Tensor.rand(20).schedule_linear().src
+    s2 = Tensor.rand(20).schedule_linear().src
 
     assert len(s) == len(s2), f"{len(s)} != {len(s2)}"
     for x,y in zip(s, s2):
-      if not (x.ast == y.ast):
-        print(f"{x.ast} != {y.ast}")
+      if not (x.src[0] == y.src[0]):
+        print(f"{x.src[0]} != {y.src[0]}")
 
     Tensor.rand(1, device=f"{Device.DEFAULT}:1").realize()
 
-    s3 = Tensor.rand(20, device=f"{Device.DEFAULT}:1").schedule()
-    s4 = Tensor.rand(20, device=f"{Device.DEFAULT}:1").schedule()
+    s3 = Tensor.rand(20, device=f"{Device.DEFAULT}:1").schedule_linear().src
+    s4 = Tensor.rand(20, device=f"{Device.DEFAULT}:1").schedule_linear().src
 
     assert len(s3) == len(s4), f"{len(s3)} != {len(s4)}"
     assert len(s2) == len(s4), f"{len(s)} != {len(s3)}"
     for x,y in zip(s3, s4):
-      if not (x.ast == y.ast):
-        print(f"{x.ast} != {y.ast}")
+      if not (x.src[0] == y.src[0]):
+        print(f"{x.src[0]} != {y.src[0]}")
 
   @unittest.skipUnless(is_dtype_supported(dtypes.bfloat16), "need bfloat16 support")
   def test_rand_bfloat16(self):

test/backend/test_renderer_failures.py

@@ -3,7 +3,8 @@ 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, get_program
+from tinygrad.engine.realize import CompiledRunner
+from tinygrad.codegen import to_program
 from tinygrad.helpers import prod
 from tinygrad.renderer.cstyle import CStyleLanguage
 from tinygrad.renderer.ptx import PTXRenderer
@@ -12,15 +13,15 @@ from tinygrad.runtime.ops_python import PythonRenderer
 from tinygrad.uop.ops import UOp, Ops, KernelInfo, python_alu
 from tinygrad.tensor import Tensor, _to_np_dtype
 
-def _test_uop_result(inputs:list[Tensor], prg, local_size=None):
+def _test_uop_result(inputs:list[Tensor], prg:UOp, local_size=None):
   for x in inputs: x.realize()
-  uops = prg.uops
+  uops = prg.src[2].src
   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]
-  prg = replace(prg, device=Device.DEFAULT)
-  if local_size is not None: prg = replace(prg, local_size=local_size)
-  ei = CompiledRunner(prg)
+  info = prg.arg
+  if local_size is not None: info = replace(info, local_size=tuple(local_size))
+  ei = CompiledRunner(prg.replace(arg=info), Device.DEFAULT)
   ei.exec(outbufs+inbufs)
   return [np.frombuffer(x.as_memoryview(), _to_np_dtype(x.dtype)) for x in outbufs]
 
@@ -33,7 +34,7 @@ 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,), arg=KernelInfo())
-  prg = get_program(sink, Device[Device.DEFAULT].renderer)
+  prg = to_program(sink, Device[Device.DEFAULT].renderer)
   return _test_uop_result([Tensor([input_val])], prg)[0]
 
 class TestRendererFailures(unittest.TestCase):
@@ -43,7 +44,7 @@ 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,), arg=KernelInfo())
-    prg = get_program(sink, Device[Device.DEFAULT].renderer)
+    prg = to_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])
 
@@ -54,7 +55,7 @@ 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,), arg=KernelInfo())
-    prg = get_program(sink, Device[Device.DEFAULT].renderer)
+    prg = to_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])
 
@@ -69,10 +70,9 @@ class TestCStyleFailures(unittest.TestCase):
     dtype = "bool" if op in (Ops.OR, Ops.XOR, Ops.AND) else None
     ret = Tensor.empty(1, dtype=dtype)
     for _ in range(5): ret = python_alu[op](ret, Tensor.empty(1, dtype=dtype))
-    schedule = ret.schedule()
-    assert len(schedule) == 1
-    schedule[0].lower()
-    src = schedule[0].prg.p.src
+    linear = ret.schedule_linear()
+    assert len(linear.src) == 1
+    src = to_program(linear.src[0].src[0], Device[Device.DEFAULT].renderer).src[3].arg
     self.assertEqual("("*5 not in src, should_strip_paren)
 
   def test_repeat_add(self): self._test_src_strip_paren(Ops.ADD)
@@ -102,7 +102,7 @@ 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,), arg=KernelInfo())
-    prg = get_program(sink, Device[Device.DEFAULT].renderer)
+    prg = to_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])
 

test/backend/test_schedule.py

@@ -12,8 +12,7 @@ from tinygrad.device import is_dtype_supported
 from tinygrad.dtype import DType
 from tinygrad.uop.ops import UOp, Ops, UPat
 from tinygrad.helpers import CI, DEBUG, OSX, GlobalCounters, Context, getenv, all_same, temp
-from tinygrad.engine.realize import CompiledRunner, run_linear
-from tinygrad.schedule import linear_to_schedule
+from tinygrad.engine.realize import compile_linear, run_linear
 
 class KernelCountException(Exception): pass
 def check_schedule(t:Tensor|list[Tensor]|UOp, allowed:int, to_prerealize:list[Tensor]|None=None, filter_sink=True):
@@ -24,17 +23,17 @@ def check_schedule(t:Tensor|list[Tensor]|UOp, allowed:int, to_prerealize:list[Te
   else:
     assert isinstance(t, UOp), f"can't schedule {t}"
     linear, var_vals = Tensor(t).linear_with_vars()
-  # test lowering all the ExecItems
-  sched = linear_to_schedule(linear)
-  for si in sched: si.lower()
-  kernel_cnt = len([si for si in sched if isinstance(si.prg, CompiledRunner) or not filter_sink])
+  kernel_cnt = sum((len(call.device) if isinstance(call.device, tuple) else 1)
+                   for call in linear.src if call.src[0].op is Ops.SINK or not filter_sink)
   if kernel_cnt != allowed:
     print(f"SCHEDULE ISSUE, expecting {allowed} got {kernel_cnt}")
     if DEBUG >= 3:
-      for i,s in enumerate(sched):
+      for i,call in enumerate(linear.src):
         print("kernel", i+1)
-        print(s.ast)
+        print(call.src[0])
     raise KernelCountException(f"{kernel_cnt} != {allowed}")
+  # test compiling the linear
+  compile_linear(linear)
   return linear, var_vals
 
 def _realize_weights(m):
@@ -50,9 +49,8 @@ def _test_conv2d(allowed:int, dtype:DType=dtypes.float):
   ret = Tensor.conv2d(img, w).relu().mean().backward()
   dtypes.default_float = old_default_float
   linear, var_vals = Tensor.linear_with_vars(ret, img.grad, w.grad)
-  s = linear_to_schedule(linear)
   run_linear(linear, var_vals)
-  cnt = len([si for si in s if si.ast.op is Ops.SINK])
+  cnt = len([call for call in linear.src if call.src[0].op is Ops.SINK])
   assert cnt == allowed, f"expected {allowed} kernels, got {cnt}"
   if getenv("CHECK", 1):
     import torch
@@ -74,7 +72,7 @@ class TestSchedule(unittest.TestCase):
     x = Tensor.arange(25).reshape(1,1,5,5).cast(dtypes.float32)
     t = x.avg_pool2d(padding=1)
     linear, var_vals = t.linear_with_vars()
-    self.assertEqual(len(linear_to_schedule(linear)), kcount)
+    self.assertEqual(len(linear.src), kcount)
     run_linear(linear, var_vals)
     import torch
     torch_out = torch.nn.functional.avg_pool2d(torch.arange(25).reshape(1,1,5,5).float(), kernel_size=(2,2), padding=1).numpy()
@@ -789,7 +787,7 @@ class TestSchedule(unittest.TestCase):
     gc.collect()
     base = GlobalCounters.mem_used
     x = Tensor.ones(256).contiguous().realize()
-    (x+Tensor.ones(256).contiguous()).schedule()
+    (x+Tensor.ones(256).contiguous()).schedule_linear()
     gc.collect()
     self.assertEqual(GlobalCounters.mem_used-base, 1024)
 
@@ -799,9 +797,8 @@ class TestSchedule(unittest.TestCase):
       def cnt():
         x, y, z = Tensor.empty((64, 64), dtype='float'), Tensor.empty((64, 64), dtype='float'), Tensor.empty((64, 64), dtype='float')
         a = (x @ y).relu()
-        sched = ((a @ z).relu() + a).schedule()
-        for si in sched: si.lower()
-        return len([si for si in sched if isinstance(si.prg, CompiledRunner)])
+        linear = compile_linear(((a @ z).relu() + a).schedule_linear())
+        return len([call for call in linear.src if call.src[0].op is Ops.PROGRAM])
 
       with Context(IMAGE=1):
         self.assertEqual(cnt(), 5)
@@ -816,9 +813,8 @@ class TestSchedule(unittest.TestCase):
         rb = (((((inp @ b1) + c1).relu() @ b2) + c2).relu() + inp).relu()
         b16, c16 = Tensor.empty((512, 16), dtype='float'), Tensor.empty((16,), dtype='float')
         b32, c32 = Tensor.empty((512, 32), dtype='float'), Tensor.empty((32,), dtype='float')
-        sched = Tensor.schedule((rb @ b16 + c16).relu(), (rb @ b32 + c32).relu())
-        for si in sched: si.lower()
-        return len([si for si in sched if isinstance(si.prg, CompiledRunner)])
+        linear = compile_linear(Tensor.schedule_linear((rb @ b16 + c16).relu(), (rb @ b32 + c32).relu()))
+        return len([call for call in linear.src if call.src[0].op is Ops.PROGRAM])
 
       with Context(IMAGE=1):
         self.assertEqual(cnt(), 9)
@@ -830,9 +826,8 @@ class TestSchedule(unittest.TestCase):
         x, y, z = Tensor.empty((1, 4, 3, 3)), Tensor.empty((4, 1, 3, 3)), Tensor.empty((4, 1, 7, 7))
         a = x.conv2d(y, Tensor.empty(4), groups=4, padding=1)
         b = a.conv2d(z, groups=4, padding=3)
-        sched = (a + b).schedule()
-        for si in sched: si.lower()
-        return len([si for si in sched if isinstance(si.prg, CompiledRunner)])
+        linear = compile_linear((a + b).schedule_linear())
+        return len([call for call in linear.src if call.src[0].op is Ops.PROGRAM])
 
       with Context(IMAGE=1):
         self.assertEqual(cnt(), 5)
@@ -1055,7 +1050,7 @@ class TestSchedule(unittest.TestCase):
     expected = (a+a2).tolist()
     a.assign(a+a2)
     linear, var_vals = a.linear_with_vars()
-    kcount = len(linear_to_schedule(linear))
+    kcount = len(linear.src)
     run_linear(linear, var_vals)
     self.assertListEqual(a.tolist(), expected)
     self.assertEqual(kcount, expected_kcount)
@@ -1334,7 +1329,7 @@ class TestCopyFolding(unittest.TestCase):
     b = Tensor.empty(4, device="CPU")
     add = a+b
     assert all_same([x.device for x in add.uop.src]), f"ALU has different devices! {[x.device for x in add.src]}"
-    add.schedule()
+    add.schedule_linear()
 
   def test_alu_before_copy(self):
     buf = Tensor.ones(1).contiguous().realize()
@@ -1356,7 +1351,7 @@ class TestCopyFolding(unittest.TestCase):
     a = Tensor.ones(4).contiguous().realize().uop.buf_uop
     t = Tensor(a.copy_to_device(a.device))
     linear, var_vals = t.linear_with_vars()
-    assert len([s for s in linear_to_schedule(linear) if s.ast.op is Ops.COPY]) == 0
+    assert len([call for call in linear.src if call.src[0].op is Ops.COPY]) == 0
     run_linear(linear, var_vals)
     assert t.uop.is_realized, f"didn't realize Tensor {t}"
     self.assertListEqual(t.tolist(), [1.,1.,1.,1.])

test/backend/test_softmax_fusion.py

@@ -2,6 +2,7 @@ import unittest
 import numpy as np
 from tinygrad import Tensor, GlobalCounters, Context, Device
 from tinygrad.dtype import DTypeLike, dtypes
+from tinygrad.engine.realize import run_linear
 from tinygrad.helpers import DEBUG, get_single_element
 from tinygrad.device import is_dtype_supported
 
@@ -26,7 +27,10 @@ def single_kernel_softmax(x_in:Tensor, axis=-1, dtype:DTypeLike|None=None) -> Te
   out = e.div(ss).reshape(x_in.shape)
   return out
 
-def run_one_schedule_item(out): get_single_element(out.schedule()).run()
+def run_one_schedule_item(out):
+  linear = out.schedule_linear()
+  get_single_element(linear.src)
+  run_linear(linear)
 
 class TestFuse(unittest.TestCase):
   def _test_fuse(self, fxn, *args, atol=1e-6, allow_multiple=False, **kwargs):
@@ -100,8 +104,8 @@ class TestFuse(unittest.TestCase):
     k = (x @ wk).contiguous()
     v = (x @ wv).contiguous()
     attn = q.scaled_dot_product_attention(k, v)
-    s = attn.schedule()
-    self.assertEqual(len(s), 4) # 3 matmul and 1 attention
+    s = attn.schedule_linear()
+    self.assertEqual(len(s.src), 4) # 3 matmul and 1 attention
 
   @unittest.skip("needs RANGEIFY>1")
   def test_flash_attention(self):

test/backend/test_uops.py

@@ -7,17 +7,16 @@ from tinygrad.dtype import dtypes, DType, AddrSpace, ConstFloat  # noqa: F401
 from tinygrad.device import Buffer, Device
 from tinygrad.uop.ops import Ops, UOp, KernelInfo, AxisType
 from tinygrad.renderer.cstyle import CStyleLanguage
-from tinygrad.engine.realize import CompiledRunner, get_program, get_runner
-from tinygrad.schedule import ExecItem
+from tinygrad.engine.realize import CompiledRunner, run_linear
+from tinygrad.codegen import to_program
 from tinygrad.device import is_dtype_supported
 from tinygrad.codegen.opt import Opt, OptOps
 from tinygrad.renderer.ptx import PTXRenderer
 from test.helpers import to_uops_list
-from dataclasses import replace
 
 def _uops_to_prg(uops_list):
-  prg = get_program(UOp.sink(*uops_list, arg=KernelInfo()), Device[Device.DEFAULT].renderer)
-  return CompiledRunner(replace(prg, device=Device.DEFAULT))
+  prg = to_program(UOp.sink(*uops_list, arg=KernelInfo()), Device[Device.DEFAULT].renderer)
+  return CompiledRunner(prg, Device.DEFAULT)
 
 def uop(uops:list[UOp], op:Ops, dtype:Optional[DType], src:tuple[UOp, ...], arg:Any=None) -> UOp:
   if op is Ops.CONST: uops.append(UOp.const(dtype, arg))
@@ -246,11 +245,11 @@ class TestAssembly(unittest.TestCase):
     a = Tensor.empty(1024)
     b = Tensor.empty(1024)
     c = (a*b).sum()
-    ast = c.schedule()[-1].ast
+    ast = c.schedule_linear().src[-1].src[0]
     opts_to_apply = [Opt(OptOps.UNROLL, 0, 4)]
     ast = ast.replace(arg=KernelInfo(opts_to_apply=tuple(opts_to_apply)))
-    program = get_program(ast, Device[Device.DEFAULT].renderer)
-    uops = program.uops
+    program = to_program(ast, Device[Device.DEFAULT].renderer)
+    uops = tuple(program.src[2].src)
     self.assertGreaterEqual(len([x.op for x in uops if x.op is Ops.MULACC]), 4)
 
   def test_mulacc_shl(self):
@@ -281,7 +280,7 @@ class TestZeroRange(unittest.TestCase):
 
 class TestUOpPrograms(unittest.TestCase):
   def _run(self, prog:UOp, *tensors:Tensor):
-    ExecItem(prog, [t.uop.buffer for t in tensors], prg=get_runner(Device.DEFAULT, prog)).run(wait=True)
+    run_linear(UOp(Ops.LINEAR, src=(prog.call(*[t.uop.buf_uop for t in tensors]),)), do_update_stats=False)
 
   def test_simple(self):
     out = Tensor.empty(10,10,dtype=dtypes.int)

test/device/test_hcq.py

@@ -6,7 +6,8 @@ from tinygrad.device import Buffer, BufferSpec
 from tinygrad.runtime.support.hcq import HCQCompiled, HCQBuffer
 from tinygrad.runtime.autogen import libc
 from tinygrad.runtime.support.system import PCIIfaceBase
-from tinygrad.engine.realize import get_runner, CompiledRunner, get_program
+from tinygrad.engine.realize import get_runner, CompiledRunner
+from tinygrad.codegen import to_program
 from tinygrad.codegen.opt import Opt, OptOps
 from tinygrad import Variable
 
@@ -19,9 +20,9 @@ class TestHCQ(unittest.TestCase):
     TestHCQ.d0 = Device[Device.DEFAULT]
     TestHCQ.a = Tensor([0.,1.], device=Device.DEFAULT).realize()
     TestHCQ.b = self.a + 1
-    si = self.b.schedule()[-1]
+    si = self.b.schedule_linear().src[-1]
 
-    TestHCQ.runner = get_runner(TestHCQ.d0.device, si.ast)
+    TestHCQ.runner = get_runner(TestHCQ.d0.device, si.src[0])
     TestHCQ.b.uop.buffer.allocate()
 
     TestHCQ.kernargs_ba_ptr = TestHCQ.runner._prg.fill_kernargs([TestHCQ.b.uop.buffer._buf, TestHCQ.a.uop.buffer._buf])
@@ -163,9 +164,10 @@ class TestHCQ(unittest.TestCase):
 
     a = Tensor.randint((3, 3, 3), dtype=dtypes.int, device=Device.DEFAULT).realize()
     b = a + 1
-    si = b.schedule()[-1]
+    si = b.schedule_linear().src[-1]
 
-    runner = CompiledRunner(get_program(replace_opts(si.ast, [Opt(op=OptOps.LOCAL, axis=0, arg=3) for _ in range(3)]), TestHCQ.d0.renderer))
+    runner = CompiledRunner(to_program(replace_opts(si.src[0], [Opt(op=OptOps.LOCAL, axis=0, arg=3) for _ in range(3)]), TestHCQ.d0.renderer),
+                            Device.DEFAULT)
 
     zb = Buffer(Device.DEFAULT, 3 * 3 * 3, dtypes.int, options=BufferSpec(cpu_access=True, nolru=True)).ensure_allocated()
     zt = Buffer(Device.DEFAULT, 3 * 3 * 3, dtypes.int, options=BufferSpec(cpu_access=True, nolru=True)).ensure_allocated()
@@ -468,7 +470,7 @@ class TestHCQ(unittest.TestCase):
   def test_memory_barrier(self):
     a = Tensor([0, 1], device=Device.DEFAULT, dtype=dtypes.int8).realize()
     b = a + 1
-    runner = get_runner(TestHCQ.d0.device, b.schedule()[-1].ast)
+    runner = get_runner(TestHCQ.d0.device, b.schedule_linear().src[-1].src[0])
 
     buf1 = Buffer(Device.DEFAULT, 2, dtypes.int8, options=BufferSpec(nolru=True)).ensure_allocated()
     buf2 = Buffer(Device.DEFAULT, 2, dtypes.int8, options=BufferSpec(cpu_access=True, nolru=True)).ensure_allocated()

test/device/test_metal.py

@@ -50,6 +50,19 @@ kernel void r_5(device int* data0, const device int* data1, uint3 gid [[threadgr
       compiled = compiled[:40] # corrupt the compiled program
       MetalProgram(device, "r_5", compiled)
 
+  def test_wait_skips_in_flight(self):
+    device = MetalDevice("metal")
+    compiled = MetalCompiler().compile("""
+#include <metal_stdlib>
+kernel void noop(uint3 gid [[threadgroup_position_in_grid]], uint3 lid [[thread_position_in_threadgroup]]) {}
+""")
+    prg = MetalProgram(device, "noop", compiled)
+    self.assertIsInstance(prg(wait=True), float)
+    self.assertEqual(device.mtl_buffers_in_flight, [])
+    self.assertIsNone(prg(wait=False))
+    self.assertEqual(len(device.mtl_buffers_in_flight), 1)
+    device.synchronize()
+
   def test_free(self):
     size = 2**16
     device = Device['METAL']

test/device/test_validate_with_cpu.py

@@ -0,0 +1,40 @@
+import unittest
+from tinygrad import Tensor, Context, Variable, Device
+from test.helpers import needs_second_gpu
+
+class TestValidateWithCPU(unittest.TestCase):
+  def setUp(self):
+    self.ctx = Context(VALIDATE_WITH_CPU=1)
+    self.ctx.__enter__()
+  def tearDown(self): self.ctx.__exit__(None, None, None)
+
+  def test_add(self): self.assertListEqual((Tensor([1.,2,3])+Tensor([4.,5,6])).tolist(), [5.0, 7.0, 9.0])
+  def test_mul(self): self.assertListEqual((Tensor([1.,2,3])*Tensor([4.,5,6])).tolist(), [4.0, 10.0, 18.0])
+  def test_sum(self): self.assertEqual(Tensor([1.,2,3,4]).sum().item(), 10.0)
+  def test_reduce_then_op(self): self.assertEqual((Tensor([1.,2,3,4]).sum() * 2).item(), 20.0)
+
+  def test_assign(self):
+    a = Tensor([1.,2,3]).realize()
+    a.assign(a + 1).realize()
+    self.assertListEqual(a.tolist(), [2.0, 3.0, 4.0])
+
+  def test_buffer_view(self):
+    self.assertListEqual((Tensor([1.,2,3,4,5,6,7,8])[2:6] + 1).tolist(), [4.0, 5.0, 6.0, 7.0])
+
+  def test_symbolic(self):
+    i = Variable('i', 1, 10)
+    ones = Tensor.ones(10).contiguous()
+    self.assertListEqual((ones[:i.bind(5)] + 1).contiguous()[:5].tolist(), [2.0]*5)
+
+  def test_multi_kernel(self):
+    a = (Tensor([1.,2,3]) + 1).contiguous()
+    b = (a * 2).contiguous()
+    self.assertListEqual((b - 1).tolist(), [3.0, 5.0, 7.0])
+
+  @needs_second_gpu
+  def test_sharded(self):
+    t = Tensor([1.,2,3,4]).shard((f"{Device.DEFAULT}:0", f"{Device.DEFAULT}:1"), axis=0)
+    self.assertListEqual((t + 1).tolist(), [2.0, 3.0, 4.0, 5.0])
+
+if __name__ == "__main__":
+  unittest.main()

test/external/external_benchmark_op_conv.py

@@ -3,7 +3,8 @@ from dataclasses import replace
 from tinygrad import dtypes, Device
 from tinygrad.uop.ops import UOp, AxisType, Ops, KernelInfo
 from tinygrad.codegen.opt import Opt, OptOps # pylint: disable=unused-import
-from tinygrad.engine.realize import CompiledRunner, get_program
+from tinygrad.engine.realize import CompiledRunner
+from tinygrad.codegen import to_program
 from tinygrad.helpers import dedup, getenv
 from tinygrad.device import Buffer
 from tinygrad.dtype import ImageDType, Invalid
@@ -88,13 +89,13 @@ ast = {143: vision_conv_143, 153: vision_conv_153, 172: dm_conv_172}[getenv("NUM
 renderer = Device.default.renderer
 allocator = Device.default.allocator
 
-ps = get_program(ast, renderer)
-cr = CompiledRunner(replace(ps, device=Device.DEFAULT))
+ps = to_program(ast, renderer)
+cr = CompiledRunner(ps, Device.DEFAULT)
 
 gs = sorted(dedup([u for u in ast.toposort() if u.op is Ops.PARAM]), key=lambda u: u.arg)
 # print(len(gs))
 # print([g.dtype for g in gs])
-bufs = [Buffer(ps.device, g.size, g.dtype if isinstance(g.dtype, ImageDType) else g.dtype._base).ensure_allocated() for g in gs]
+bufs = [Buffer(ps.arg.device, g.size, g.dtype if isinstance(g.dtype, ImageDType) else g.dtype._base).ensure_allocated() for g in gs]
 
 t = cr(bufs, wait=True)
 print(f"{t*1e6:.2f} us")

test/external/external_benchmark_schedule.py

@@ -23,10 +23,10 @@ if __name__ == "__main__":
     if not FORWARD_ONLY:
       with Timing("***** model schedule in  "):
         with Profiling(PROFILE >= 3):
-          sched = out.schedule()
+          linear = out.schedule_linear()
 
       if not SCHEDULE_ONLY:
-        asts = list({x.ast.key:x.ast for x in sched if x.ast.op is Ops.SINK}.values())
+        asts = list({call.src[0].key:call.src[0] for call in linear.src if call.src[0].op is Ops.SINK}.values())
         if (restrict_kernel := getenv("RESTRICT_KERNEL", -1)) != -1: asts = asts[restrict_kernel:restrict_kernel+1]
 
         with Profiling(PROFILE, fn="/tmp/rewrite.prof"):

test/external/external_test_hcq.py

@@ -20,8 +20,8 @@ class TestHCQ(unittest.TestCase):
     #TestHCQ.d1: AMDDevice = Device["AMD:1"]
     TestHCQ.a = Tensor([0.,1.], device=Device.DEFAULT).realize()
     TestHCQ.b = self.a + 1
-    si = self.b.schedule()[-1]
-    TestHCQ.runner = get_runner(TestHCQ.d0.device, si.ast)
+    linear = self.b.schedule_linear()
+    TestHCQ.runner = get_runner(TestHCQ.d0.device, linear.src[-1].src[0])
     TestHCQ.b.uop.buffer.allocate()
     # wow that's a lot of abstraction layers
     TestHCQ.addr = struct.pack("QQ", TestHCQ.b.uop.buffer._buf, TestHCQ.a.uop.buffer._buf)

test/external/external_test_onnx_runner.py

@@ -10,8 +10,8 @@ from hypothesis import given, strategies as st
 # copied from test_const_folding.py
 def _check_ast_count(desired_count:int, t:Tensor):
   # NOTE: this has side effect because everything can be scheduled only once
-  schedule = t.schedule()
-  asts = [s for s in schedule if s.ast.op is Ops.SINK]
+  linear = t.schedule_linear()
+  asts = [call for call in linear.src if call.src[0].op is Ops.SINK]
   assert len(asts) == desired_count, f"{len(asts)} != {desired_count}"
 
 def build_onnx(nodes, from_disk:bool=True, **kwargs):

test/external/external_test_opt.py

@@ -7,7 +7,6 @@ from tinygrad import GlobalCounters, Tensor, Device
 from tinygrad.helpers import getenv
 from tinygrad.nn.state import get_parameters
 from tinygrad.engine.realize import capturing, run_linear
-from tinygrad.schedule import linear_to_schedule
 from tinygrad.tensor import _to_np_dtype
 
 class CLCache:
@@ -15,7 +14,7 @@ class CLCache:
     self.allowed, self.strict, self.preclear, self.var_vals = allowed, strict, preclear, var_vals if var_vals is not None else {}
     self.count = 0
   def add_linear(self, linear, var_vals):
-    self.count += len(linear_to_schedule(linear))
+    self.count += len(linear.src)
     run_linear(linear, var_vals)
   def __enter__(self):
     if self.preclear:

test/external/external_test_schedule_scaling.py

@@ -6,7 +6,7 @@ class TestScheduleScaling(unittest.TestCase):
 
   def _assert_linear(self, fn, n_small=200, n_large=1000):
     """Assert schedule time scales at most ~linearly: time(n_large)/time(n_small) should be close to n_large/n_small."""
-    fn(n_small).schedule()  # warmup
+    fn(n_small).schedule_linear()  # warmup
     t_small = min(self._time_schedule(fn, n) for n in [n_small]*3)
     t_large = min(self._time_schedule(fn, n) for n in [n_large]*3)
     size_ratio = n_large / n_small  # 5.0
@@ -19,7 +19,7 @@ class TestScheduleScaling(unittest.TestCase):
   @staticmethod
   def _time_schedule(fn, n) -> float:
     st = time.perf_counter()
-    fn(n).schedule()
+    fn(n).schedule_linear()
     return time.perf_counter() - st
 
   # *** rangeify: ending_ranges accumulation and consumer merge ***

test/external/external_uop_gc.py

@@ -1,7 +1,8 @@
 import gc
 from tinygrad import Tensor, UOp, Device, nn
 from tinygrad.schedule import schedule_cache
-from tinygrad.engine.realize import method_cache, get_program
+from tinygrad.engine.realize import method_cache
+from tinygrad.codegen import to_program, to_program_cache
 from tinygrad.schedule.indexing import apply_movement_op, _apply_reshape
 from tinygrad.uop.divandmod import fold_divmod_general
 from test.test_tiny import TestTiny
@@ -14,13 +15,13 @@ def print_uops():
 def start(): pass
 def single_tensor(): Tensor([2])
 def two_plus_two(): Tensor([2])+Tensor([2])
-def two_plus_two_schedule(): (Tensor([2])+Tensor([2])).schedule()
+def two_plus_two_schedule(): (Tensor([2])+Tensor([2])).schedule_linear()
 def two_plus_two_kernel():
-  si = (Tensor([2])+Tensor([2])).schedule()[-1]
-  get_program(si.ast, Device.default.renderer)
+  linear = (Tensor([2])+Tensor([2])).schedule_linear()
+  to_program(linear.src[-1].src[0], Device.default.renderer)
 def two_plus_two_linearize():
-  si = (Tensor([2])+Tensor([2])).schedule()[-1]
-  get_program(si.ast, Device.default.renderer)
+  linear = (Tensor([2])+Tensor([2])).schedule_linear()
+  to_program(linear.src[-1].src[0], Device.default.renderer)
 def two_plus_two_realize(): (Tensor([2])+Tensor([2])).realize()
 def two_plus_two_item(): (Tensor([2])+Tensor([2])).item()
 def gradient_test():
@@ -36,8 +37,8 @@ def kernel_matmul():
   x = Tensor.eye(3, requires_grad=True)
   y = Tensor([[2.0,0,-2.0]], requires_grad=True)
   z = y.matmul(x)
-  si = z.schedule()[-1]
-  get_program(si.ast, Device.default.renderer)
+  linear = z.schedule_linear()
+  to_program(linear.src[-1].src[0], Device.default.renderer)
 def realized_matmul():
   x = Tensor.eye(3, requires_grad=True)
   y = Tensor([[2.0,0,-2.0]], requires_grad=True)
@@ -71,6 +72,7 @@ if __name__ == "__main__":
     # these caches will keep uops alive
     schedule_cache.clear()
     method_cache.clear()
+    to_program_cache.clear()
     apply_movement_op.cache_clear()
     _apply_reshape.cache_clear()
     fold_divmod_general.cache_clear()

test/external/fuzz_graph.py

@@ -4,8 +4,7 @@ from tinygrad.device import Buffer, Device
 from tinygrad.helpers import Context, getenv, from_mv
 from tinygrad.dtype import dtypes
 from tinygrad.tensor import Tensor, _to_np_dtype
-from tinygrad.engine.realize import BufferXfer, get_runner
-from tinygrad.schedule import ExecItem
+from tinygrad.engine.realize import BufferXfer, get_runner, ExecItem
 from tinygrad.uop.ops import UOp, Ops
 from tinygrad.engine.jit import apply_graph_to_jit
 
@@ -20,8 +19,8 @@ def gen_prg(device, inputs_cnt):
     s = fst[0]
     for i in range(1, inputs_cnt): s = s.bitwise_xor(fst[i])
 
-    si = s.schedule()[-1]
-    prg = get_runner(device, si.ast)
+    linear = s.schedule_linear()
+    prg = get_runner(device, linear.src[-1].src[0])
   cached_prgs[(device, inputs_cnt)] = prg
   return prg
 

test/helpers.py

@@ -5,7 +5,7 @@ import numpy as np
 from tinygrad import Tensor, dtypes, Device
 from tinygrad.uop.ops import UOp, Ops, KernelInfo
 from tinygrad.tensor import _to_np_dtype
-from tinygrad.engine.realize import get_program
+from tinygrad.codegen import to_program
 from tinygrad.dtype import DType
 from tinygrad.nn.state import get_parameters
 from tinygrad.helpers import T, CI, Target
@@ -80,8 +80,8 @@ def eval_uop(uop:UOp, inputs:list[tuple[DType, list[Any]]]|None=None):
     bufs.append(buf:=allocator.alloc(len(data) * buf_dt.itemsize))
     allocator._copyin(buf, memoryview(struct.pack(str(len(data)) + (buf_dt.fmt or ""), *data)))
   g = UOp(Ops.PARAM, uop.dtype.ptr(), arg=0, src=())
-  prg = get_program(UOp.store(g.index(UOp.const(dtypes.int, 0)), uop).sink(arg=KernelInfo()), PythonRenderer(Target("PYTHON")))
-  prog = PythonProgram("run", PythonCompiler().compile(prg.src))
+  prg = to_program(UOp.store(g.index(UOp.const(dtypes.int, 0)), uop).sink(arg=KernelInfo()), PythonRenderer(Target("PYTHON")))
+  prog = PythonProgram("run", PythonCompiler().compile(prg.src[3].arg))
   prog(out_buf:=allocator.alloc(uop.dtype.itemsize), *bufs)
   return out_buf.cast(uop.dtype.fmt or "").tolist()[0]
 

test/mockgpu/usb.py

@@ -202,6 +202,8 @@ class MockASM24State:
     return None
 
 class MockUSB3:
+  @classmethod
+  def list_devices(cls, vendor, dev): return [(0, "usb:mock")]
   def __init__(self, *args, **kwargs):
     self.product, self.is_custom = "", False
   def send_batch(self, cdbs:list[bytes], idata:list[int]|None=None, odata:list[bytes|None]|None=None) -> list[bytes|None]:

test/null/test_attention.py

@@ -16,9 +16,9 @@ class TestAttention(unittest.TestCase):
     k = Tensor.ones(BS, seqlen, dim, dtype=dtypes.half).contiguous().realize()
     v = Tensor.ones(BS, seqlen, dim, dtype=dtypes.half).contiguous().realize()
     attn = q.scaled_dot_product_attention(k, v)
-    sched = attn.schedule()
+    sched = attn.schedule_linear()
     # attention has 4 kernels now
-    self.assertEqual(len(sched), 4)
+    self.assertEqual(len(sched.src), 4)
 
   def test_apply_rope_jit_prune(self):
     def rope_fn(x_in, pos): return apply_rope(x_in, pos)

test/null/test_compile_failures.py

@@ -3,11 +3,12 @@ from contextlib import redirect_stdout
 from tinygrad import Tensor, dtypes, Device
 from tinygrad.helpers import OSX, DEV
 from tinygrad.device import is_dtype_supported
-from tinygrad.engine.realize import get_program
+from tinygrad.engine.realize import compile_linear
+from tinygrad.codegen import to_program
 
 class TestCompileFailures(unittest.TestCase):
   def compile(self, out:Tensor):
-    for si in out.schedule(): si.lower()
+    compile_linear(out.schedule_linear())
 
   @unittest.skipUnless(is_dtype_supported(dtypes.uchar), f"no uint8 on {Device.DEFAULT}")
   def test_interpolate_atari(self):
@@ -21,9 +22,9 @@ class TestDisassembly(unittest.TestCase):
   @unittest.skipUnless(Device.DEFAULT in ("CPU",) and DEV.renderer not in ("LLVM", "LVP") and OSX, "m series cpus support fp16 arithmetic")
   def test_float16_alu(self):
     c = Tensor([1], dtype=dtypes.float16) + Tensor([1], dtype=dtypes.float16)
-    s = c.schedule()[-1]
-    p = get_program(s.ast, Device[Device.DEFAULT].renderer)
-    lib = Device[Device.DEFAULT].compiler.compile(p.src)
+    s = c.schedule_linear().src[-1]
+    p = to_program(s.src[0], Device[Device.DEFAULT].renderer)
+    lib = Device[Device.DEFAULT].compiler.compile(p.src[3].arg)
     out = io.StringIO()
     with redirect_stdout(out): Device[Device.DEFAULT].compiler.disassemble(lib)
     assert "fcvt" not in out.getvalue()

test/null/test_const_folding.py

@@ -7,8 +7,8 @@ import numpy as np
 
 def _check_ast_count(desired_count:int, t:Tensor):
   # NOTE: this has side effect because everything can be scheduled only once
-  schedule = t.schedule()
-  asts = [s for s in schedule if s.ast.op is Ops.SINK]
+  linear = t.schedule_linear()
+  asts = [s for s in linear.src if s.src[0].op is Ops.SINK]
   len(asts)
   # NOT SUPPORTED ANYMORE
   #assert len(asts) == desired_count, f"{len(asts)} != {desired_count}"

test/null/test_device.py

@@ -153,7 +153,7 @@ class TestDevVar(unittest.TestCase):
       self.assertEqual(DEV.target("CPU"), Target("CPU"))
 
   def test_dev_arch_override(self):
-    with Context(DEV="NULL:HIP:gfx1100"):
+    with Context(DEV="NULL::gfx1100"):
       self.assertEqual(Device["NULL"].renderer.target.arch, "gfx1100")
 
 class MockCompiler(Compiler):

test/null/test_gc.py

@@ -60,7 +60,7 @@ class TestGC(unittest.TestCase):
     init = bufs_allocated()
     x = Tensor.ones(256).contiguous().realize()
     y = Tensor.ones(5, 5).contiguous()
-    y.schedule()
+    y.schedule_linear()
     del x
     del y
     self.assertEqual(bufs_allocated()-init, 0)

test/null/test_linearizer_failures.py

@@ -2,8 +2,8 @@
 import unittest
 from tinygrad.uop.ops import UOp, Ops, AxisType, KernelInfo
 from tinygrad.dtype import dtypes
-from tinygrad.engine.realize import get_program
 from tinygrad.device import Device
+from tinygrad.codegen import to_program
 
 class TestLinearizerFailures(unittest.TestCase):
   def test_fail_1(self):
@@ -19,7 +19,7 @@ class TestLinearizerFailures(unittest.TestCase):
     c9 = ((((c6+(c8*UOp.const(dtypes.float, -1.0)))*(c6+(c8*UOp.const(dtypes.float, -1.0)))).reduce(c5, arg=Ops.ADD)*UOp.const(dtypes.float, 0.000390625))+UOp.const(dtypes.float, 1e-05)).sqrt().reciprocal()
     c10 = c0.index(c3).store(c9).end(c1, c2)
     ast = c10.sink(arg=KernelInfo())
-    get_program(ast, renderer=Device[Device.DEFAULT].renderer)
+    to_program(ast, renderer=Device[Device.DEFAULT].renderer)
 
 if __name__ == '__main__':
   unittest.main()

test/null/test_linearizer_rewrite.py

@@ -1,6 +1,6 @@
 import unittest
 from tinygrad import Tensor, Context, Device
-from tinygrad.engine.realize import get_program
+from tinygrad.codegen import to_program
 from tinygrad.codegen.opt import Opt, OptOps
 from tinygrad.uop.ops import KernelInfo
 
@@ -9,37 +9,37 @@ class TestLinearizerRewrite(unittest.TestCase):
     t = Tensor.ones((64,64), device="NULL").contiguous().realize()
     out = (t*2).sum(axis=1)
     with Context(SPLIT_REDUCEOP=0, DEVECTORIZE=0):
-      si = out.schedule()[-1]
+      si = out.schedule_linear().src[-1]
       opts_to_apply = []
       opts_to_apply.append(Opt(OptOps.UPCAST, 0, 4))
       opts_to_apply.append(Opt(OptOps.UNROLL, 0, 4))
-      ast = si.ast.replace(arg=KernelInfo(opts_to_apply=tuple(opts_to_apply)))
-      prg = get_program(ast, Device["CPU"].renderer)
-      print(prg.src)
+      ast = si.src[0].replace(arg=KernelInfo(opts_to_apply=tuple(opts_to_apply)))
+      prg = to_program(ast, Device["CPU"].renderer)
+      print(prg.src[3].arg)
 
   def test_arange(self):
     out = Tensor.arange(32, device="NULL")
     with Context(SPLIT_REDUCEOP=0, DEVECTORIZE=0):
-      si = out.schedule()[-1]
+      si = out.schedule_linear().src[-1]
       opts_to_apply = []
       opts_to_apply.append(Opt(OptOps.UPCAST, 0, 4))
-      ast = si.ast.replace(arg=KernelInfo(opts_to_apply=tuple(opts_to_apply)))
-      prg = get_program(ast, Device["CPU"].renderer)
-      print(prg.src)
+      ast = si.src[0].replace(arg=KernelInfo(opts_to_apply=tuple(opts_to_apply)))
+      prg = to_program(ast, Device["CPU"].renderer)
+      print(prg.src[3].arg)
 
   def test_kernel_info(self):
     out = Tensor.arange(4, device="NULL")
-    si = out.schedule()[-1]
+    si = out.schedule_linear().src[-1]
 
-    ast = si.ast.replace(arg=KernelInfo(opts_to_apply=()))
-    prg = get_program(ast, Device["CPU"].renderer)
-    assert prg.applied_opts == (), f"expected no opts, got {prg}"
+    ast = si.src[0].replace(arg=KernelInfo(opts_to_apply=()))
+    prg = to_program(ast, Device["CPU"].renderer)
+    assert prg.src[0].arg.applied_opts == (), f"expected no opts, got {prg}"
 
-    prg = get_program(ast.replace(arg=KernelInfo()), Device["CPU"].renderer)
-    assert prg.applied_opts != (), f"expected opts to apply, got {prg.applied_opts}"
+    prg = to_program(ast.replace(arg=KernelInfo()), Device["CPU"].renderer)
+    assert prg.src[0].arg.applied_opts != (), f"expected opts to apply, got {prg.src[0].arg.applied_opts}"
 
-    prg = get_program(ast.replace(arg=KernelInfo(name="custom")), Device["CPU"].renderer)
-    self.assertEqual(prg.name, "custom")
+    prg = to_program(ast.replace(arg=KernelInfo(name="custom")), Device["CPU"].renderer)
+    self.assertEqual(prg.arg.name, "custom")
 
 if __name__ == '__main__':
   unittest.main()

test/null/test_process_replay.py

@@ -9,7 +9,7 @@ N = 16
 class TestProcessReplay(unittest.TestCase):
   @classmethod
   def setUpClass(cls):
-    cls.ast = (Tensor.empty(N, N) @ Tensor.empty(N, N)).schedule()[-1].ast
+    cls.ast = (Tensor.empty(N, N) @ Tensor.empty(N, N)).schedule_linear().src[-1].src[0]
     cls.renderer = Device[Device.DEFAULT].renderer
 
   def test_replay_no_opts(self):
@@ -35,9 +35,9 @@ class TestProcessReplay(unittest.TestCase):
 
   def test_beam(self):
     with Context(BEAM=1):
-      si = (Tensor.empty(N, N) @ Tensor.empty(N, N)).schedule()[-1]
-    p = do_to_program(si.ast, self.renderer)
-    good, compare, _ = replay_to_program(p, si.ast, self.renderer)
+      ast = (Tensor.empty(N, N) @ Tensor.empty(N, N)).schedule_linear().src[-1].src[0]
+    p = do_to_program(ast, self.renderer)
+    good, compare, _ = replay_to_program(p, ast, self.renderer)
     self.assertEqual(good, compare)
 
 if __name__ == '__main__':

test/null/test_schedule.py

@@ -3,8 +3,8 @@ import gc, unittest, time
 from tinygrad import nn, dtypes, Device, Tensor
 from tinygrad.uop.ops import UOp, Ops, GroupOp, UPat, KernelInfo
 from tinygrad.helpers import DEBUG, GlobalCounters, Context
-from tinygrad.engine.realize import CompiledRunner, run_linear
-from tinygrad.schedule import linear_to_schedule
+from tinygrad.engine.realize import compile_linear, run_linear
+from tinygrad.codegen import to_program
 
 class KernelCountException(Exception): pass
 def check_schedule(t:Tensor|list[Tensor]|UOp, allowed:int, to_prerealize:list[Tensor]|None=None, filter_sink=True):
@@ -15,17 +15,17 @@ def check_schedule(t:Tensor|list[Tensor]|UOp, allowed:int, to_prerealize:list[Te
   else:
     assert isinstance(t, UOp), f"can't schedule {t}"
     linear, var_vals = Tensor(t).linear_with_vars()
-  # test lowering all the ExecItems
-  sched = linear_to_schedule(linear)
-  for si in sched: si.lower()
-  kernel_cnt = len([si for si in sched if isinstance(si.prg, CompiledRunner) or not filter_sink])
+  kernel_cnt = sum((len(call.device) if isinstance(call.device, tuple) else 1)
+                   for call in linear.src if call.src[0].op is Ops.SINK or not filter_sink)
   if kernel_cnt != allowed:
     print(f"SCHEDULE ISSUE, expecting {allowed} got {kernel_cnt}")
     if DEBUG >= 3:
-      for i,s in enumerate(sched):
+      for i,call in enumerate(linear.src):
         print("kernel", i+1)
-        print(s.ast)
+        print(call.src[0])
     raise KernelCountException(f"{kernel_cnt} != {allowed}")
+  # test compiling the linear
+  compile_linear(linear)
   return linear, var_vals
 
 def _realize_weights(m):
@@ -88,7 +88,7 @@ class TestBufferUOp(unittest.TestCase):
     # unused variable should not appear in var_vals even when there's other work
     a = Tensor(UOp.variable("unused", 0, 10).bind(1))
     b = Tensor.empty(3) + 1
-    _, var_vals = Tensor.schedule_with_vars(a, b)
+    _, var_vals = Tensor.linear_with_vars(a, b)
     self.assertEqual(var_vals, {})
     self.assertIsNone(a.uop.base.realized)
 
@@ -142,7 +142,7 @@ class TestSimpleSchedule(unittest.TestCase):
     a = Tensor.empty(16,16).sum(axis=1)
     a1 = a.reshape(4,4)
     a2 = a.reshape(16,1,1)
-    self.assertEqual(len(Tensor.schedule(a1, a2)), 1)
+    self.assertEqual(len(Tensor.schedule_linear(a1, a2).src), 1)
 
 class TestSchedule(unittest.TestCase):
   def test_create_schedule_handles_multi_kernel_after_and_after_deps(self):
@@ -167,8 +167,8 @@ class TestSchedule(unittest.TestCase):
     kc = Tensor.custom_kernel(out, src_after, fxn=named_copy("kc"))[0]
     out_after = Tensor(kc.uop.src[0].after(*kc.uop.src[1:], kd.uop))
 
-    schedule = out_after.schedule()
-    names = [si.ast.arg.name for si in schedule]
+    linear = out_after.schedule_linear()
+    names = [call.src[0].arg.name for call in linear.src]
     self.assertEqual(set(names), {"ka", "kb", "kc", "kd"})
     self.assertEqual(names[-1], "kc")
     self.assertLess(names.index("ka"), names.index("kc"))
@@ -209,10 +209,10 @@ class TestSchedule(unittest.TestCase):
     t = Tensor.zeros((3, 3)).contiguous().realize()
     v = t[1]  # view - is_realized but not has_buffer_identity
     assert v.uop.is_realized
-    sched, _ = Tensor.schedule_with_vars(v)
-    self.assertEqual(len(sched), 0)
+    linear, _ = Tensor.linear_with_vars(v)
+    self.assertEqual(len(linear.src), 0)
 
-  # NOTE: because empty does not have a lowered ExecItem if realize is called on a childless empty, it never gets allocated.
+  # NOTE: because empty does not have a lowered kernel if realize is called on a childless empty, it never gets allocated.
   def test_childless_empty_never_allocates(self):
     a = Tensor.empty(10)
     a.realize()
@@ -668,9 +668,9 @@ class TestSchedule(unittest.TestCase):
     check_schedule(c, 2)
 
   def _alu_from_tensor(self, t:Tensor):
-    s = [s for s in t.schedule() if s.ast.op is Ops.SINK]
+    s = [s for s in t.schedule_linear().src if s.src[0].op is Ops.SINK]
     self.assertEqual(len(s), 1)
-    return [u.op for u in s[0].ast.toposort() if u.op in GroupOp.ALU]
+    return [u.op for u in s[0].src[0].toposort() if u.op in GroupOp.ALU]
 
   def test_2_pow_is_exp2(self):
     t = 2.0 ** Tensor([1.0, 2.0, 3.0])
@@ -799,12 +799,12 @@ class TestSchedule(unittest.TestCase):
     Tensor.manual_seed(0)
     x = Tensor.randn(4, 12, 64, 64, dtype=dtypes.half).realize()
     out = x.softmax(dtype=dtypes.float)
-    sched = out.schedule()
-    self.assertEqual(len(sched), 3)
+    linear = out.schedule_linear()
+    self.assertEqual(len(linear.src), 3)
     # max reduction stays in input dtype (no numerical loss), upcast happens after subtracting max
-    self.assertEqual(sched[0].bufs[0].dtype, dtypes.half)
-    self.assertEqual(sched[1].bufs[0].dtype, dtypes.float)
-    self.assertEqual(sched[2].bufs[0].dtype, dtypes.float)
+    self.assertEqual(linear.src[0].src[1].dtype, dtypes.half)
+    self.assertEqual(linear.src[1].src[1].dtype, dtypes.float)
+    self.assertEqual(linear.src[2].src[1].dtype, dtypes.float)
 
   def test_softmax_backward(self):
     Tensor.manual_seed(0)
@@ -961,7 +961,7 @@ class TestSchedule(unittest.TestCase):
     gc.collect()
     base = GlobalCounters.mem_used
     Tensor.ones(256).contiguous().realize()
-    Tensor.ones(5, 5).contiguous().schedule()
+    Tensor.ones(5, 5).contiguous().schedule_linear()
     gc.collect()
     self.assertEqual(GlobalCounters.mem_used-base, 0)
 
@@ -1174,24 +1174,24 @@ class TestFusionOp(unittest.TestCase):
     st = time.perf_counter()
     a = Tensor([1,2,3,4])
     for _ in range(24): a = a + a
-    sched = a.schedule()
-    sched[-1].lower()
+    linear = a.schedule_linear()
+    prg = to_program(linear.src[-1].src[0], renderer=Device[Device.DEFAULT].renderer)
     self.assertLess(time.perf_counter()-st, 2.0)
-    assert len(sched[-1].prg.p.src.splitlines()) < 250
+    assert len(prg.src[3].arg.splitlines()) < 250
 
   def test_recursive_add_cmp(self):
     st = time.perf_counter()
     a = Tensor([1,2,3,4])
     for _ in range(24): a = a + a
-    sched1 = a.schedule()
+    linear1 = a.schedule_linear()
     b = Tensor([1,2,3,4])
     for _ in range(24): b = b + b
-    sched2 = b.schedule()
+    linear2 = b.schedule_linear()
     c = Tensor([1,2,3,4])
     for _ in range(23): c = c + c
-    sched3 = c.schedule()
-    self.assertEqual(sched1[-1].ast, sched2[-1].ast)
-    with self.assertRaises(AssertionError): self.assertEqual(sched1[-1].ast, sched3[-1].ast)
+    linear3 = c.schedule_linear()
+    self.assertEqual(linear1.src[-1].src[0], linear2.src[-1].src[0])
+    with self.assertRaises(AssertionError): self.assertEqual(linear1.src[-1].src[0], linear3.src[-1].src[0])
     self.assertLess(time.perf_counter()-st, 2.0)
 
   def test_recursive_pad(self):
@@ -1199,8 +1199,8 @@ class TestFusionOp(unittest.TestCase):
     val = 1.0
     a = Tensor(val)
     for _ in range(24): a = Tensor.stack(a, a)[0]
-    sched = a.schedule()
-    self.assertLessEqual(len(sched), 1)
+    linear = a.schedule_linear()
+    self.assertLessEqual(len(linear.src), 1)
     self.assertLess(time.perf_counter()-st, 2.0)
 
   def test_recursive_reshape(self):
@@ -1209,8 +1209,8 @@ class TestFusionOp(unittest.TestCase):
     b = Tensor.empty(16, 2).realize()
     r = a.sum(1)
     for _ in range(24): r = r.reshape(16, 2) + b
-    sched = r.schedule()
-    self.assertEqual(len(sched), 1)
+    linear = r.schedule_linear()
+    self.assertEqual(len(linear.src), 1)
     self.assertLess(time.perf_counter()-st, 2.0)
 
 # NOTE: the NULL backend supports BUFFER_VIEW

test/null/test_schedule_cache.py

@@ -4,7 +4,7 @@ from tinygrad.helpers import cpu_events
 from tinygrad.schedule import schedule_cache
 
 def schedule_one():
-  Tensor([1]).schedule()
+  Tensor([1]).schedule_linear()
 
 class TestScheduleCache(unittest.TestCase):
   def test_bound_variable_var_vals(self):
@@ -12,7 +12,7 @@ class TestScheduleCache(unittest.TestCase):
     x = Tensor.ones(10).contiguous().realize()
 
     t = x + Tensor(v.bind(42))
-    _, var_vals = t.schedule_with_vars()
+    _, var_vals = t.linear_with_vars()
     self.assertEqual(var_vals, {'pos': 42})
 
   def test_disable_schedule_cache(self):

test/null/test_tensor.py

@@ -6,7 +6,7 @@ from tinygrad.device import is_dtype_supported
 from tinygrad.uop.ops import Ops, UOp
 from tinygrad.renderer.ptx import PTXRenderer
 from tinygrad.renderer.nir import NIRRenderer
-from tinygrad.engine.realize import get_program
+from tinygrad.codegen import to_program
 from tinygrad.dtype import DType
 
 x_init = np.random.randn(1,3).astype(np.float32)
@@ -62,12 +62,13 @@ class TestIdxUpcast(unittest.TestCase):
     for src in ast.src:
       if (ret:=self._find_op(src, op)) is not None: return ret
   def _schedule_render(self, a: Tensor):
-    schedule, _ = a.schedule_with_vars()
-    for s in schedule:
-      if s.ast.op is Ops.SINK:
-        renderer = Device[s.bufs[0].device].renderer
-        prg = get_program(s.ast, renderer)
-        return prg.uops
+    linear, _ = a.linear_with_vars()
+    for si in linear.src:
+      ast = si.src[0]
+      if ast.op is Ops.SINK:
+        renderer = Device[si.src[1].buffer.device].renderer
+        prg = to_program(ast, renderer)
+        return tuple(prg.src[2].src)
 
   def _assert(self, dtype: DType, a: Tensor):
     uops = self._schedule_render(a)
@@ -162,9 +163,9 @@ class TestRand(unittest.TestCase):
   def test_rand_large_tensor(self):
     # large tensor rand (num > uint32.max) should not crash in frontend
     Tensor.manual_seed(0)
-    Tensor.rand(2**17, 2**17).schedule()
-    Tensor.rand(2**17, 2**17).schedule()
-    Tensor.rand(2**17, 2**17).schedule()
+    Tensor.rand(2**17, 2**17).schedule_linear()
+    Tensor.rand(2**17, 2**17).schedule_linear()
+    Tensor.rand(2**17, 2**17).schedule_linear()
 
 class TestTensorConstLike(unittest.TestCase):
   def test_const_like_shape(self):

test/null/test_tensor_metadata.py

@@ -2,7 +2,6 @@ import unittest
 from tinygrad import Tensor, dtypes
 from tinygrad.tensor import _METADATA
 from tinygrad.engine.realize import capturing
-from tinygrad.schedule import linear_to_schedule
 from tinygrad.helpers import Context
 
 @unittest.skip("tensor metadata is no longer supported")
@@ -18,41 +17,41 @@ class TestTensorMetadata(unittest.TestCase):
   def test_exclude_noop_metadata(self):
     a = Tensor.rand(4, 4)*1
     self.assertEqual(a.uop.metadata[0].name, "__mul__")
-    k = a.schedule()[-1]
-    self.assertEqual([m.name for m in k.metadata], ["rand"])
+    k = a.schedule_linear().src[-1]
+    self.assertEqual([m.name for m in k.arg.metadata], ["rand"])
 
   @unittest.skip("metadata not reaching kernel schedule")
   def test_exclude_const_metadata(self):
     a = Tensor.arange(4)
     b = Tensor.full((4,), -1, dtype=dtypes.int).contiguous()
-    sched = Tensor.schedule(a, b)
-    self.assertEqual([m.name for m in sched[0].metadata], ["arange"])
-    self.assertEqual([m.name for m in sched[1].metadata], ["contiguous"])
+    sched = a.schedule_linear(b)
+    self.assertEqual([m.name for m in sched.src[0].arg.metadata], ["arange"])
+    self.assertEqual([m.name for m in sched.src[1].arg.metadata], ["contiguous"])
 
   def test_matmul(self):
     x = Tensor.rand(3, requires_grad=True)
     W = Tensor.rand(3, 3, requires_grad=True)
     out = x.matmul(W)
     self.assertEqual(out.uop.metadata[0].name, "matmul")
-    si = out.schedule()[-1]
-    self.assertEqual(len(si.metadata), 1)
-    self.assertEqual(si.metadata[0].name, "matmul")
+    si = out.schedule_linear().src[-1]
+    self.assertEqual(len(si.arg.metadata), 1)
+    self.assertEqual(si.arg.metadata[0].name, "matmul")
 
   def test_relu(self):
     x = Tensor.rand(3, requires_grad=True)
     out = x.relu()
     self.assertEqual(out.uop.metadata[0].name, "relu")
-    si = out.schedule()[-1]
-    self.assertEqual(len(si.metadata), 1)
-    self.assertEqual(si.metadata[0].name, "relu")
+    si = out.schedule_linear().src[-1]
+    self.assertEqual(len(si.arg.metadata), 1)
+    self.assertEqual(si.arg.metadata[0].name, "relu")
 
   @unittest.skip("assign metadata no longer captured")
   def test_assign(self):
     x = Tensor.empty(10, 10).realize()
     x.assign(Tensor.ones(10, 10).contiguous())
-    si = x.schedule()[-1]
-    self.assertEqual(len(si.metadata), 1)
-    self.assertEqual(si.metadata[0].name, "assign")
+    si = x.schedule_linear().src[-1]
+    self.assertEqual(len(si.arg.metadata), 1)
+    self.assertEqual(si.arg.metadata[0].name, "assign")
 
   def test_complex(self):
     x = Tensor.rand(3, requires_grad=True)
@@ -61,9 +60,9 @@ class TestTensorMetadata(unittest.TestCase):
     self.assertEqual(out.uop.metadata[0].name, "__mul__")
     self.assertEqual(out.uop.src[0].metadata[0].name, "relu")
     self.assertEqual(out.uop.src[1].metadata[0].name, "sigmoid")
-    si = out.schedule()[-1]
-    self.assertEqual(len(si.metadata), 3)
-    self.assertEqual(set(m.name for m in si.metadata), {"relu", "sigmoid", "__mul__"})
+    si = out.schedule_linear().src[-1]
+    self.assertEqual(len(si.arg.metadata), 3)
+    self.assertEqual(set(m.name for m in si.arg.metadata), {"relu", "sigmoid", "__mul__"})
 
   @unittest.skip("flaky")
   def test_complex_backward(self):
@@ -76,10 +75,10 @@ class TestTensorMetadata(unittest.TestCase):
     #self.assertTrue(x.grad.uop.metadata[0].backward)  # TODO: backward flag is False
     self.assertEqual(y.grad.uop.metadata[0].name, "sigmoid")
     #self.assertTrue(y.grad.uop.metadata[0].backward)  # TODO: backward flag is False
-    si = Tensor.schedule(out, x.grad, y.grad)[-1]
-    #self.assertEqual(len(si.metadata), 3, f"failed with {si.metadata}")
+    si = out.schedule_linear(x.grad, y.grad).src[-1]
+    #self.assertEqual(len(si.arg.metadata), 3, f"failed with {si.arg.metadata}")
     # skip numpy, this is schedule cache
-    self.assertSetEqual(set(m.name for m in si.metadata if m.name != "numpy"), {"sigmoid", "relu"})
+    self.assertSetEqual(set(m.name for m in si.arg.metadata if m.name != "numpy"), {"sigmoid", "relu"})
     #bw = [m for m in si.metadata if m.backward]
     #self.assertEqual(len(bw), 1)
     #self.assertEqual(bw[0].name, "sigmoid")
@@ -91,16 +90,16 @@ class TestTensorMetadata(unittest.TestCase):
       out = (x.relu() * y.sigmoid()).sum()
       self.assertIsNone(out.uop.metadata)
       self.assertIsNone(out.uop.src[0].metadata)
-      si = out.schedule()[-1]
-      self.assertEqual(si.metadata, ())
+      si = out.schedule_linear().src[-1]
+      self.assertEqual(si.arg.metadata, ())
 
   def _has_metadata(self, h, name):
     linears = []
     capturing.append(type("", (), {"add_linear": lambda _, linear, var_vals: linears.append(linear)})())
     try: h.realize()
     finally: capturing.clear()
-    items = [ei for linear in linears for ei in linear_to_schedule(linear)]
-    return any(m.name == name for ei in items for m in ei.metadata)
+    calls = [call for linear in linears for call in linear.src]
+    return any(m.name == name for call in calls for m in call.arg.metadata)
 
   def test_metadata_survives_realize_pending_assign(self):
     shared = Tensor.rand(4)

test/null/test_tensor_uop_mixin.py

@@ -112,6 +112,31 @@ class TestTensorUOpOneHot(unittest.TestCase):
     t = _t(5)
     self.assertIs(_strip_unique(t.one_hot(5).uop), _strip_unique(t.uop.one_hot(5)))
 
+class TestTensorUOpSort(unittest.TestCase):
+  def _check(self, t, **kw):
+    tv, ti = t.sort(**kw)
+    uv, ui = t.uop.sort(**kw)
+    self.assertIs(_strip_unique(tv.uop), _strip_unique(uv))
+    self.assertIs(_strip_unique(ti.uop), _strip_unique(ui))
+  def test_sort_1d(self):         self._check(Tensor([0.5, 0.1, 0.3]).float())
+  def test_sort_descending(self): self._check(Tensor([0.5, 0.1, 0.3]).float(), descending=True)
+  def test_sort_2d(self):         self._check(_t(2, 4).float())
+  def test_sort_single(self):     self._check(Tensor([1.0]).float())
+  def test_argsort(self):
+    t = Tensor([0.5, 0.1, 0.3]).float()
+    self.assertIs(_strip_unique(t.argsort().uop), _strip_unique(t.uop.argsort()))
+  def test_topk(self):
+    t = _t(2, 4).float()
+    tv, ti = t.topk(2)
+    uv, ui = t.uop.topk(2)
+    self.assertIs(_strip_unique(tv.uop), _strip_unique(uv))
+    self.assertIs(_strip_unique(ti.uop), _strip_unique(ui))
+
+class TestTensorUOpAllclose(unittest.TestCase):
+  def test_allclose(self):
+    a, b = _t(4).float(), _t(4).float()
+    self.assertIs(_strip_unique(a.allclose(b).uop), _strip_unique(a.uop.allclose(b.uop)))
+
 class TestTensorUOpGather(unittest.TestCase):
   def _check(self, t, dim, idx):
     self.assertIs(_strip_unique(t.gather(dim, idx).uop), _strip_unique(t.uop.gather(dim, idx.uop)))
@@ -138,6 +163,48 @@ class TestTensorUOpLoss(unittest.TestCase):
     t, Y = _t(2, 3).float(), Tensor([1, 2], dtype=dtypes.int32)
     self.assertIs(_strip_unique(t.sparse_categorical_crossentropy(Y, ignore_index=0).uop),
                   _strip_unique(t.uop.sparse_categorical_crossentropy(Y.uop, ignore_index=0)))
+  def test_nll_loss(self):
+    t, Y = _t(2, 3).float().log_softmax(), Tensor([1, 2], dtype=dtypes.int32)
+    self.assertIs(_strip_unique(t.nll_loss(Y).uop), _strip_unique(t.uop.nll_loss(Y.uop)))
+  def test_nll_loss_weight(self):
+    t, Y, w = _t(2, 3).float().log_softmax(), Tensor([1, 2], dtype=dtypes.int32), _t(3).float()
+    self.assertIs(_strip_unique(t.nll_loss(Y, weight=w).uop), _strip_unique(t.uop.nll_loss(Y.uop, weight=w.uop)))
+  def test_nll_loss_ignore_index(self):
+    t, Y = _t(2, 3).float().log_softmax(), Tensor([1, 2], dtype=dtypes.int32)
+    self.assertIs(_strip_unique(t.nll_loss(Y, ignore_index=1).uop), _strip_unique(t.uop.nll_loss(Y.uop, ignore_index=1)))
+  def test_nll_loss_none_reduction(self):
+    t, Y = _t(2, 3).float().log_softmax(), Tensor([1, 2], dtype=dtypes.int32)
+    self.assertIs(_strip_unique(t.nll_loss(Y, reduction="none").uop), _strip_unique(t.uop.nll_loss(Y.uop, reduction="none")))
+  def test_nll_loss_weight_ignore_index(self):
+    t, Y, w = _t(2, 3).float().log_softmax(), Tensor([1, 2], dtype=dtypes.int32), _t(3).float()
+    self.assertIs(_strip_unique(t.nll_loss(Y, weight=w, ignore_index=1).uop),
+                  _strip_unique(t.uop.nll_loss(Y.uop, weight=w.uop, ignore_index=1)))
+
+class TestTensorUOpScatter(unittest.TestCase):
+  def test_scatter(self):
+    x, idx, src = _t(3, 4).float(), Tensor([[0, 1, 2, 0]], dtype=dtypes.int32), _t(1, 4).float()
+    self.assertIs(_strip_unique(x.scatter(0, idx, src).uop), _strip_unique(x.uop.scatter(0, idx.uop, src.uop)))
+  def test_scatter_scalar_src(self):
+    x, idx = _t(3, 4).float(), Tensor([[0, 1]], dtype=dtypes.int32)
+    self.assertIs(_strip_unique(x.scatter(1, idx, 3.14).uop), _strip_unique(x.uop.scatter(1, idx.uop, 3.14)))
+  # inf cannot be cast to int — this regresses if scalar src is routed through index.dtype first
+  def test_scatter_inf_src(self):
+    x, idx = _t(3, 4).float(), Tensor([[0, 1]], dtype=dtypes.int32)
+    self.assertIs(_strip_unique(x.scatter(1, idx, float("inf")).uop),
+                  _strip_unique(x.uop.scatter(1, idx.uop, float("inf"))))
+  def test_scatter_add(self):
+    x, idx = _t(3, 4).float(), Tensor([[0, 1]], dtype=dtypes.int32)
+    self.assertIs(_strip_unique(x.scatter(1, idx, 3.14, reduce="add").uop),
+                  _strip_unique(x.uop.scatter(1, idx.uop, 3.14, reduce="add")))
+  def test_scatter_multiply(self):
+    x, idx = _t(3, 4).float(), Tensor([[0, 1]], dtype=dtypes.int32)
+    self.assertIs(_strip_unique(x.scatter(1, idx, 3.14, reduce="multiply").uop),
+                  _strip_unique(x.uop.scatter(1, idx.uop, 3.14, reduce="multiply")))
+  # tensor src with reduce hits the "elif reduce: raise" branch in both Tensor and UOp paths
+  def test_scatter_tensor_src_with_reduce_raises(self):
+    x, idx, src = _t(3, 4).float(), Tensor([[0, 1]], dtype=dtypes.int32), _t(1, 2).float()
+    with self.assertRaises(TypeError): x.scatter(1, idx, src, reduce="add")
+    with self.assertRaises(TypeError): x.uop.scatter(1, idx.uop, src.uop, reduce="add")
 
 class TestTensorUOpScatterReduce(unittest.TestCase):
   def _check(self, x, idx, src, **kw):
@@ -176,6 +243,17 @@ class TestTensorUOpCat(unittest.TestCase):
   def test_cat_3tensors(self): _check(self, _t(2, 3), lambda x: x.cat(x, x, dim=0))
   def test_cat_neg_dim(self):  _check(self, _t(2, 3, 4), lambda x: x.cat(x, dim=-1))
 
+class TestTensorUOpPad(unittest.TestCase):
+  def test_pad_flat(self):               _check(self, _t(4, 5), lambda x: x.pad((1, 2, 0, 3)))
+  def test_pad_flat_negative(self):      _check(self, _t(4, 5), lambda x: x.pad((1, -1, 0, 2), value=-1.0))
+  def test_pad_grouped_none(self):       _check(self, _t(4, 5), lambda x: x.pad((None, (0, 3))))
+  def test_pad_circular(self):           _check(self, _t(4, 5), lambda x: x.pad(((1, 2), (0, 3)), mode="circular"))
+  def test_pad_circular_zero_after(self):_check(self, _t(4, 5), lambda x: x.pad(((1, 0), (2, 0)), mode="circular"))
+  def test_pad_reflect(self):            _check(self, _t(4, 5), lambda x: x.pad(((1, 2), (0, 3)), mode="reflect"))
+  def test_pad_reflect_negative(self):   _check(self, _t(4, 5), lambda x: x.pad(((1, -1), (0, 2)), mode="reflect"))
+  def test_pad_replicate(self):          _check(self, _t(4, 5), lambda x: x.pad(((1, 2), (0, 3)), mode="replicate"))
+  def test_pad_replicate_negative(self): _check(self, _t(4, 5), lambda x: x.pad(((1, -1), (0, 2)), mode="replicate"))
+
 class TestTensorUOpStack(unittest.TestCase):
   def test_stack_dim0(self):     _check(self, _t(2, 3), lambda x: x.stack(x, dim=0))
   def test_stack_dim1(self):     _check(self, _t(2, 3), lambda x: x.stack(x, dim=1))

test/null/test_tensor_uop_representation.py

@@ -16,7 +16,7 @@ class TestTensorMutates(unittest.TestCase):
     pa = a.uop
     pb = b.uop
     pr = ret.uop
-    ret.schedule()
+    ret.schedule_linear()
     self.assertIsNot(pa, a.uop)
     self.assertIsNot(pb, b.uop)
     self.assertIsNot(pr, ret.uop)

test/null/test_tinyfs.py

@@ -5,22 +5,22 @@ class TestLoadStore(unittest.TestCase):
   def test_load_shape(self):
     t = Tensor(bytes(16)).fs_load(1024)
     assert t.shape == (1024,), t.shape
-    t.schedule()
+    t.schedule_linear()
 
   def test_store_shape(self):
     t = Tensor.zeros(1024).fs_store()
     assert t.shape == (16,), t.shape
-    t.schedule()
+    t.schedule_linear()
 
   def test_load_large_shape(self):
     t = Tensor(bytes(16)).fs_load(10_000_000)
     assert t.shape == (10_000_000,), t.shape
-    t.schedule()
+    t.schedule_linear()
 
   def test_store_large_shape(self):
     t = Tensor.zeros(10_000_000).fs_store()
     assert t.shape == (16,), t.shape
-    t.schedule()
+    t.schedule_linear()
 
 if __name__ == "__main__":
   unittest.main()

test/null/test_uops.py

@@ -228,7 +228,7 @@ class TestUOpMethod(unittest.TestCase):
     a = UOp.variable("a", 1, 10)
     uop_var = Tensor(a.bind(1))
     st_var = Tensor.empty((2, 10))[:, :a.bind(1)]
-    _, var_vals = (uop_var+st_var).schedule_with_vars()
+    _, var_vals = (uop_var+st_var).linear_with_vars()
     self.assertEqual(len(var_vals), 1)
     self.assertEqual(list(var_vals)[0], a.expr)
 

test/null/test_uops_stats.py

@@ -1,8 +1,8 @@
 import unittest
 from tinygrad import Tensor
 from tinygrad.helpers import GlobalCounters, DEV
-from tinygrad.engine.realize import get_program
-from tinygrad.renderer import ProgramSpec
+from tinygrad.engine.realize import compile_linear, estimate_uop
+from tinygrad.codegen import to_program
 from tinygrad.renderer import Estimates
 from tinygrad.uop.ops import Ops, UOp
 from tinygrad.dtype import dtypes
@@ -18,8 +18,8 @@ def flops_mem(uops, ignore_indexing=False):
 # **************** new FlopCounter ****************
 
 def get_stats(x:Tensor):
-  si = x.schedule()[-1].lower()
-  return si.prg.estimates.ops, si.prg.estimates.mem
+  est = estimate_uop(compile_linear(x.schedule_linear()).src[-1])
+  return est.ops, est.mem
 
 @unittest.skipIf(Device.DEFAULT == "WEBGPU", "webgpu does extra load/store for packed types")
 class TestMemoryCount(unittest.TestCase):
@@ -165,75 +165,75 @@ N = 64
 class TestStatsOptimized(unittest.TestCase):
   @classmethod
   def setUpClass(cls):
-    cls.ast_gemm = (Tensor.empty(N, N) @ Tensor.empty(N, N)).schedule()[-1].ast
-    cls.ast_reduce = (Tensor.empty(N*N).sum()).schedule()[-1].ast
+    cls.ast_gemm = (Tensor.empty(N, N) @ Tensor.empty(N, N)).schedule_linear().src[-1].src[0]
+    cls.ast_reduce = (Tensor.empty(N*N).sum()).schedule_linear().src[-1].src[0]
 
-  def check_gemm(self, p:ProgramSpec, extra_flops=0):
-    #p.uops.print()
-    #print(p.src)
-    print(p.name, p.estimates.ops, p.estimates.mem, p.estimates.lds)
-    self.assertEqual(p.estimates.ops, 2*N*N*N + extra_flops)  # N**3 mulaccs
-    self.assertEqual(p.estimates.mem, 3*N*N*4) # 3 NxN mats with floats
+  def check_gemm(self, p:UOp, extra_flops=0):
+    est = p.src[0].arg.estimates
+    print(p.arg.name, est.ops, est.mem, est.lds)
+    self.assertEqual(est.ops, 2*N*N*N + extra_flops)  # N**3 mulaccs
+    self.assertEqual(est.mem, 3*N*N*4) # 3 NxN mats with floats
 
   def test_gemm(self):
-    p = get_program(replace_opts(self.ast_gemm, []), renderer=Device[Device.DEFAULT].renderer)
+    p = to_program(replace_opts(self.ast_gemm, []), renderer=Device[Device.DEFAULT].renderer)
     self.check_gemm(p)
-    self.assertEqual(p.estimates.lds, 2*N*N*N*4 + 4*N*N)
+    self.assertEqual(p.src[0].arg.estimates.lds, 2*N*N*N*4 + 4*N*N)
 
   def test_gemm_tc_unroll(self):
     try:
-      p = get_program(replace_opts(self.ast_gemm, [Opt(OptOps.TC, 0, (-1, 0, 1)), Opt(OptOps.UNROLL, 0, 2)]),
+      p = to_program(replace_opts(self.ast_gemm, [Opt(OptOps.TC, 0, (-1, 0, 1)), Opt(OptOps.UNROLL, 0, 2)]),
                       renderer=Device[Device.DEFAULT].renderer)
     except KernelOptError:
       raise unittest.SkipTest("no tensor cores")
-    print(p.src)
+    print(p.src[3].arg)
     self.check_gemm(p)
 
   # this is a good lesson about why UPCASTing is a good idea
 
   def test_gemm_one_upcasted(self):
-    p = get_program(replace_opts(self.ast_gemm, [Opt(OptOps.UPCAST, 0, 4)]), renderer=Device[Device.DEFAULT].renderer)
+    p = to_program(replace_opts(self.ast_gemm, [Opt(OptOps.UPCAST, 0, 4)]), renderer=Device[Device.DEFAULT].renderer)
     self.check_gemm(p)
-    self.assertEqual(p.estimates.lds, N*N*N*4 + N*N*N*4//4 + 4*N*N)
+    self.assertEqual(p.src[0].arg.estimates.lds, N*N*N*4 + N*N*N*4//4 + 4*N*N)
 
   def test_gemm_upcasted(self):
-    p = get_program(replace_opts(self.ast_gemm, [Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UPCAST, 1, 4), Opt(OptOps.UNROLL, 0, 4)]),
+    p = to_program(replace_opts(self.ast_gemm, [Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UPCAST, 1, 4), Opt(OptOps.UNROLL, 0, 4)]),
                     renderer=Device[Device.DEFAULT].renderer)
     self.check_gemm(p)
-    self.assertEqual(p.estimates.lds, 2*N*N*N*4//4 + 4*N*N)
+    self.assertEqual(p.src[0].arg.estimates.lds, 2*N*N*N*4//4 + 4*N*N)
 
   def test_gemm_upcasted_locals(self):
     try:
-      p = get_program(replace_opts(self.ast_gemm, [Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UPCAST, 1, 4), Opt(OptOps.LOCAL, 0, 4),
+      p = to_program(replace_opts(self.ast_gemm, [Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UPCAST, 1, 4), Opt(OptOps.LOCAL, 0, 4),
                                                    Opt(OptOps.LOCAL, 1, 4)]), renderer=Device[Device.DEFAULT].renderer)
     except KernelOptError:
       raise unittest.SkipTest("no locals")
     self.check_gemm(p)
-    self.assertEqual(p.estimates.lds, 2*N*N*N*4//4 + 4*N*N)
+    self.assertEqual(p.src[0].arg.estimates.lds, 2*N*N*N*4//4 + 4*N*N)
 
   def test_gemm_group(self):
     try:
-      p = get_program(replace_opts(self.ast_gemm, [Opt(OptOps.GROUP, 0, 4)]), renderer=Device[Device.DEFAULT].renderer)
+      p = to_program(replace_opts(self.ast_gemm, [Opt(OptOps.GROUP, 0, 4)]), renderer=Device[Device.DEFAULT].renderer)
     except KernelOptError:
       raise unittest.SkipTest("no locals")
     SZ = N*N*4
     # NOTE: these are sort of wrong. they aren't honoring the IF statement
     self.check_gemm(p, extra_flops=SZ*4)
-    self.assertEqual(p.estimates.lds, 2*N*N*N*4 + SZ*4 + (SZ*4 + 4*N*N)*4)
+    self.assertEqual(p.src[0].arg.estimates.lds, 2*N*N*N*4 + SZ*4 + (SZ*4 + 4*N*N)*4)
 
   def test_reduce(self):
-    p = get_program(replace_opts(self.ast_reduce, []), renderer=Device[Device.DEFAULT].renderer)
-    print(p.name, p.estimates.ops, p.estimates.mem, p.estimates.lds)
-    self.assertEqual(p.estimates.ops, N*N)
-    self.assertEqual(p.estimates.mem, N*N*4 + 4)
+    p = to_program(replace_opts(self.ast_reduce, []), renderer=Device[Device.DEFAULT].renderer)
+    est = p.src[0].arg.estimates
+    print(p.arg.name, est.ops, est.mem, est.lds)
+    self.assertEqual(est.ops, N*N)
+    self.assertEqual(est.mem, N*N*4 + 4)
 
   def test_reduce_group(self):
     try:
-      p = get_program(replace_opts(self.ast_reduce, [Opt(OptOps.GROUP, 0, 50)]), renderer=Device[Device.DEFAULT].renderer)
+      p = to_program(replace_opts(self.ast_reduce, [Opt(OptOps.GROUP, 0, 50)]), renderer=Device[Device.DEFAULT].renderer)
     except KernelOptError:
       raise unittest.SkipTest("no locals")
-    # NOTE: these are wrong, they don't respect the if statement
-    print(p.name, p.estimates.ops, p.estimates.mem, p.estimates.lds)
+    est = p.src[0].arg.estimates
+    print(p.arg.name, est.ops, est.mem, est.lds)
 
 if __name__ == '__main__':
   unittest.main(verbosity=2)

test/null/test_viz.py

@@ -13,6 +13,7 @@ from tinygrad.device import Buffer
 from tinygrad.uop.ops import tracked_keys, tracked_ctxs, uop_fields, active_rewrites, active_group, _name_cnt, RewriteTrace
 from tinygrad.viz.serve import load_rewrites, get_full_rewrite, uop_to_json, VizData
 from tinygrad.codegen import to_program_cache
+from tinygrad.codegen import to_program
 
 @track_rewrites(name=True)
 def exec_rewrite(sink:UOp, pm_lst:list[PatternMatcher], names:None|list[str]=None) -> UOp:
@@ -320,27 +321,27 @@ class TestVizGC(unittest.TestCase):
 # VIZ integrates with other parts of tinygrad
 
 from tinygrad import Tensor, Device
-from tinygrad.engine.realize import get_program
+from tinygrad.engine.realize import get_runner
 
 class TestVizIntegration(unittest.TestCase):
   # codegen supports rendering of code blocks
   def test_codegen_tracing(self):
     with save_viz() as viz:
-      ast = Tensor.schedule(Tensor.empty(4)+Tensor.empty(4))[0].ast
-      prg = get_program(ast, Device[Device.DEFAULT].renderer)
+      ast = (Tensor.empty(4)+Tensor.empty(4)).schedule_linear().src[0].src[0]
+      prg = to_program(ast, Device[Device.DEFAULT].renderer)
     lst = viz.list_items()
     self.assertEqual(len(lst), 3)
     self.assertEqual(lst[0]["name"], "Callify 1 Buffer n1")
     self.assertEqual(lst[1]["name"], "Schedule 1 Kernel n1")
-    self.assertEqual(lst[2]["name"], prg.name)
+    self.assertEqual(lst[2]["name"], prg.arg.name)
 
   # schedule graph CALL nodes have a link to jump to codegen
   def test_link_sched_codegen(self):
     with save_viz() as viz:
       c1 = Tensor.empty(4).add(1)
       c2 = Tensor.empty(8).add(1)
-      sched = Tensor.schedule(c1, c2)
-      prgs = [get_program(si.ast, Device[Device.DEFAULT].renderer).name for si in sched]
+      sched = c1.schedule_linear(c2)
+      prgs = [to_program(si.src[0], Device[Device.DEFAULT].renderer).arg.name for si in sched.src]
     lst = viz.list_items()
     sched_idx = next(i for i,l in enumerate(lst) if l["name"].startswith("Schedule"))
     viz_kernel = next(i for i,s in enumerate(lst[sched_idx]["steps"]) if s["name"] == "View Kernel Graph")
@@ -356,7 +357,7 @@ class TestVizIntegration(unittest.TestCase):
       a = Tensor.empty(1)
       b = Tensor.empty(1)
       metadata = (alu:=a+b).uop.metadata
-      alu.schedule()
+      alu.schedule_linear()
     graph = next(viz.get_details(0, 0))["graph"]
     self.assertEqual(len([n for n in graph.values() if repr(metadata) in n["label"]]), 1)
 
@@ -722,10 +723,10 @@ class TestCfg(unittest.TestCase):
       gidx = UOp.special(1, "gidx0")
       sink = UOp.sink(out.base, lidx, gidx, arg=KernelInfo(name=name))
       return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg="NULL"), UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))
-    with Context(DEV=f"NULL:HIP:{self.arch}"):
+    with Context(DEV=f"NULL::{self.arch}"):
       out = Tensor.custom_kernel(Tensor.empty(1), fxn=fxn)[0]
-      prg = out.schedule()[-1].lower().prg.p
-      return amdgpu_cfg(prg.lib, self.arch)
+      runner = get_runner(out.device, out.schedule_linear().src[-1].src[0])
+      return amdgpu_cfg(runner.prg.src[4].arg, self.arch)
 
   def test_simple(self):
     k = Kernel(arch=self.arch)
@@ -924,7 +925,7 @@ class TestCLI(unittest.TestCase):
       (p:=Path(tmpdir)/"profile.pkl").write_bytes(pickle.dumps(cpu_events))
       # reconstruct DEBUG=4 output and see all markers.
       with Context(DEBUG=4):
-        kernels = run_cli("--rewrites-path", str(r), "--profile-path", str(p), "-p", "-s", "NULL")
+        kernels = run_cli("--rewrites-path", str(r), "--profile-path", str(p), "-s", "NULL")
       self.assertIn("void custom_empty_n0", kernels)
       self.assertIn("marker @ 1", kernels)
       self.assertIn("void custom_empty_n1", kernels)
@@ -933,11 +934,11 @@ class TestCLI(unittest.TestCase):
       self.assertIn("UOp.const", kernels)
       # get the top slowest functions across all devices
       with Context(DEBUG=2):
-        times = run_cli("--rewrites-path", str(r), "--profile-path", str(p), "-p", "-s", "ALL", "--top", "-1")
+        times = run_cli("--rewrites-path", str(r), "--profile-path", str(p), "-s", "ALL", "--top", "-1")
       self.assertIn("TINY", times)
       self.assertIn("NULL", times)
       with Context(DEBUG=3):
-        json_lines = run_cli("--rewrites-path", str(r), "--profile-path", str(p), "-p", "-s", "ALL", "--json")
+        json_lines = run_cli("--rewrites-path", str(r), "--profile-path", str(p), "-s", "ALL", "--json")
       for line in json_lines.split("\n"): _ = json.loads(line)
 
 if __name__ == "__main__":

test/null/test_winograd.py

@@ -18,14 +18,14 @@ class TestWinograd(unittest.TestCase):
   def test_forward_kernels(self):
     x,w = Tensor.rand(1,4,9,9).realize(), Tensor.rand(4,4,3,3).realize()
     out = Tensor.conv2d(x,w)
-    self.assertEqual(len(out.schedule()), 2)
+    self.assertEqual(len(out.schedule_linear().src), 2)
 
   def test_backward_kernels(self):
     x,w = Tensor.empty(1,4,9,9,requires_grad=True).realize(), Tensor.empty(4,4,3,3,requires_grad=True).realize()
     out = Tensor.conv2d(x,w, padding=1)
     out.mean().backward()
-    backward_schedule = Tensor.schedule(x.grad, w.grad)
-    self.assertEqual(len(backward_schedule), 4)
+    backward_schedule = x.grad.schedule_linear(w.grad)
+    self.assertEqual(len(backward_schedule.src), 4)
 
   def test_counters(self):
     IC, OC, X, Y = 4,4,9,9

test/opt/test_gen_float4.py

@@ -1,8 +1,9 @@
 import unittest
 from tinygrad import Device, Tensor, dtypes
 from tinygrad.uop.ops import UOp, Ops
+from tinygrad.codegen import to_program
 from tinygrad.codegen.opt import Opt, OptOps
-from tinygrad.engine.realize import get_program
+
 from tinygrad.helpers import DEV
 from test.helpers import replace_opts
 
@@ -24,12 +25,12 @@ class TestFloat4(unittest.TestCase):
     b = Tensor.empty(2, 8).realize()
     c = a + b
 
-    s = c.schedule()[0]
-    realized_ast = s.ast
+    s = c.schedule_linear().src[0]
+    realized_ast = s.src[0]
     opts_to_apply = [Opt(op=OptOps.UPCAST, axis=0, arg=4)]
-    program = get_program(replace_opts(realized_ast, opts_to_apply), renderer=Device[Device.DEFAULT].renderer)
+    program = to_program(replace_opts(realized_ast, opts_to_apply), renderer=Device[Device.DEFAULT].renderer)
 
-    assert TestFloat4.count_float4(program.uops) == (2, 1)
+    assert TestFloat4.count_float4(tuple(program.src[2].src)) == (2, 1)
 
   @unittest.skipIf(Device.DEFAULT in {"CPU"} and AMX, "CPU with AMX upcasts float up to size 16")
   def test_float4_multidim(self):
@@ -37,9 +38,9 @@ class TestFloat4(unittest.TestCase):
     b = Tensor.empty(2, 8).realize()
     c = a + b
 
-    s = c.schedule()[0]
-    uops = get_program(replace_opts(s.ast, [Opt(op=OptOps.UPCAST, axis=0, arg=4), Opt(op=OptOps.UPCAST, axis=0, arg=2)]),
-                       renderer=Device[Device.DEFAULT].renderer).uops
+    s = c.schedule_linear().src[0]
+    uops = tuple(to_program(replace_opts(s.src[0], [Opt(op=OptOps.UPCAST, axis=0, arg=4), Opt(op=OptOps.UPCAST, axis=0, arg=2)]),
+                       renderer=Device[Device.DEFAULT].renderer).src[2].src)
     assert TestFloat4.count_float4(uops) == (4, 2)
 
   @unittest.skipUnless(Device.DEFAULT in {"CPU"} and AMX, "Only CPU with AMX upcasts float up to size 16")
@@ -49,9 +50,9 @@ class TestFloat4(unittest.TestCase):
       b = Tensor.empty(2, size).realize()
       c = a + b
 
-      s = c.schedule()[0]
-      return get_program(replace_opts(s.ast, [Opt(op=OptOps.UPCAST, axis=0, arg=4), Opt(op=OptOps.UPCAST, axis=0, arg=shift)]),
-                         renderer=Device[Device.DEFAULT].renderer).uops
+      s = c.schedule_linear().src[0]
+      return tuple(to_program(replace_opts(s.src[0], [Opt(op=OptOps.UPCAST, axis=0, arg=4), Opt(op=OptOps.UPCAST, axis=0, arg=shift)]),
+                         renderer=Device[Device.DEFAULT].renderer).src[2].src)
 
     sizes = [12, 8, 16]
     shifts = [3, 2, 4]
@@ -66,12 +67,12 @@ class TestFloat4(unittest.TestCase):
     b = Tensor.empty(9).realize().shrink(((1, 9),))
     c = a + b
 
-    s = c.schedule()[0]
-    realized_ast = s.ast
+    s = c.schedule_linear().src[0]
+    realized_ast = s.src[0]
     opts_to_apply = [Opt(op=OptOps.UPCAST, axis=0, arg=4)]
-    program = get_program(replace_opts(realized_ast, opts_to_apply), renderer=Device[Device.DEFAULT].renderer)
+    program = to_program(replace_opts(realized_ast, opts_to_apply), renderer=Device[Device.DEFAULT].renderer)
 
-    assert TestFloat4.count_float4(program.uops) == (0, 1)
+    assert TestFloat4.count_float4(tuple(program.src[2].src)) == (0, 1)
 
   @unittest.skipIf(Device.DEFAULT in {"CPU"} and AMX, "CPU with AMX upcasts float up to size 16")
   def test_float4_multidim_unaligned_load(self):
@@ -79,9 +80,9 @@ class TestFloat4(unittest.TestCase):
     b = Tensor.empty(2, 9).realize().shrink(((0, 2), (1, 9),))
     c = a + b
 
-    s = c.schedule()[0]
-    uops = get_program(replace_opts(s.ast, [Opt(op=OptOps.UPCAST, axis=1, arg=4), Opt(op=OptOps.UPCAST, axis=1, arg=2)]),
-                       renderer=Device[Device.DEFAULT].renderer).uops
+    s = c.schedule_linear().src[0]
+    uops = tuple(to_program(replace_opts(s.src[0], [Opt(op=OptOps.UPCAST, axis=1, arg=4), Opt(op=OptOps.UPCAST, axis=1, arg=2)]),
+                       renderer=Device[Device.DEFAULT].renderer).src[2].src)
 
     assert TestFloat4.count_float4(uops) == (0, 2)
 
@@ -92,9 +93,9 @@ class TestFloat4(unittest.TestCase):
       b = Tensor.empty(2, size).realize().shrink(((0, 2), (1, size),))
       c = a + b
 
-      s = c.schedule()[0]
-      return get_program(replace_opts(s.ast, [Opt(op=OptOps.UPCAST, axis=1, arg=4), Opt(op=OptOps.UPCAST, axis=1, arg=shift)]),
-                         renderer=Device[Device.DEFAULT].renderer).uops
+      s = c.schedule_linear().src[0]
+      return tuple(to_program(replace_opts(s.src[0], [Opt(op=OptOps.UPCAST, axis=1, arg=4), Opt(op=OptOps.UPCAST, axis=1, arg=shift)]),
+                         renderer=Device[Device.DEFAULT].renderer).src[2].src)
 
     sizes = [13, 9, 17]
     shifts = [3, 2, 4]
@@ -111,8 +112,8 @@ class TestFloat4(unittest.TestCase):
     # only the first and last conv dot products are aligned in a, and b is never aligned, so no
     # float4 should be emitted (the reduce axis of size 4 is the float4 axis here)
 
-    s = c.schedule()[0]
-    uops = get_program(replace_opts(s.ast, [Opt(op=OptOps.UNROLL, axis=0, arg=4)]), renderer=Device[Device.DEFAULT].renderer).uops
+    s = c.schedule_linear().src[0]
+    uops = tuple(to_program(replace_opts(s.src[0], [Opt(op=OptOps.UNROLL, axis=0, arg=4)]), renderer=Device[Device.DEFAULT].renderer).src[2].src)
 
     assert TestFloat4.count_float4(uops) == (0, 0)
 
@@ -125,9 +126,9 @@ class TestFloat4(unittest.TestCase):
     # don't.
     # UPDATE: now we do this fusion
 
-    s = c.schedule()[0]
-    uops = get_program(replace_opts(s.ast, [Opt(op=OptOps.UPCAST, axis=0, arg=0), Opt(op=OptOps.UNROLL, axis=0, arg=0)]),
-                       renderer=Device[Device.DEFAULT].renderer).uops
+    s = c.schedule_linear().src[0]
+    uops = tuple(to_program(replace_opts(s.src[0], [Opt(op=OptOps.UPCAST, axis=0, arg=0), Opt(op=OptOps.UNROLL, axis=0, arg=0)]),
+                       renderer=Device[Device.DEFAULT].renderer).src[2].src)
 
     assert TestFloat4.count_float4(uops) in {(0,1), (1,1)}
 
@@ -139,8 +140,8 @@ class TestFloat4(unittest.TestCase):
     # we will upcast the top axis of sz 4. they should not be coalesced into float4,
     # since the top axis is not contiguous.
 
-    s = c.schedule()[0]
-    uops = get_program(replace_opts(s.ast, [Opt(op=OptOps.UPCAST, axis=0, arg=4)]), renderer=Device[Device.DEFAULT].renderer).uops
+    s = c.schedule_linear().src[0]
+    uops = tuple(to_program(replace_opts(s.src[0], [Opt(op=OptOps.UPCAST, axis=0, arg=4)]), renderer=Device[Device.DEFAULT].renderer).src[2].src)
 
     assert TestFloat4.count_float4(uops) == (0, 1)
 
@@ -151,8 +152,8 @@ class TestFloat4(unittest.TestCase):
 
     # should float4 b but not a
 
-    s = c.schedule()[0]
-    uops = get_program(replace_opts(s.ast, [Opt(op=OptOps.UPCAST, axis=0, arg=4)]), renderer=Device[Device.DEFAULT].renderer).uops
+    s = c.schedule_linear().src[0]
+    uops = tuple(to_program(replace_opts(s.src[0], [Opt(op=OptOps.UPCAST, axis=0, arg=4)]), renderer=Device[Device.DEFAULT].renderer).src[2].src)
 
     assert TestFloat4.count_float4(uops) == (1, 1)
 

test/opt/test_tensor_cores.py

@@ -1,6 +1,5 @@
 import numpy as np
 import unittest
-from dataclasses import replace
 
 from tinygrad import Device, Tensor, dtypes
 from tinygrad.tensor import _to_np_dtype
@@ -9,14 +8,15 @@ from tinygrad.dtype import DType
 from tinygrad.device import is_dtype_supported
 from tinygrad.helpers import DEV, Context
 from test.helpers import slow, replace_opts
-from tinygrad.engine.realize import CompiledRunner, get_program
+from tinygrad.engine.realize import CompiledRunner
+from tinygrad.codegen import to_program
 from tinygrad.codegen.opt import Opt, OptOps, KernelOptError
 from tinygrad.codegen.opt.tc import amd_cdna_1616128
 
 # TODO: write a clean version of this
 from test.backend.test_linearizer import helper_realized_ast, helper_linearizer_opt
 
-# NOTE: get_program always passes in Device[Device.DEFAULT].renderer explicitly for process_replay!!!
+# NOTE: to_program always passes in Device[Device.DEFAULT].renderer explicitly for process_replay!!!
 
 AMX = "AMX" in DEV.arch
 
@@ -24,19 +24,19 @@ def helper_tc_ensure_uops_and_opts_count(N: int, M:int, K:int, dtype_in:DType, d
                                          ensure_triggered:bool=True):
   a, b = Tensor.rand(M, K, dtype=dtype_in), Tensor.rand(K, N, dtype=dtype_in)
   r = a.matmul(b, dtype=dtype_out)
-  sched = r.schedule()
-  realized_ast = sched[-1].ast
+  sched = r.schedule_linear()
+  realized_ast = sched.src[-1].src[0]
   opts_to_apply = [Opt(OptOps.TC, axis, (tc_select, tc_opt, 1))]
 
   if ensure_triggered:
-    program = get_program(replace_opts(realized_ast, opts_to_apply), Device[Device.DEFAULT].renderer)
-    wmmas = len([uop for uop in program.uops if uop.op is Ops.WMMA])
-    tcs = len([x for x in program.applied_opts if x.op is OptOps.TC])
+    program = to_program(replace_opts(realized_ast, opts_to_apply), Device[Device.DEFAULT].renderer)
+    wmmas = len([uop for uop in tuple(program.src[2].src) if uop.op is Ops.WMMA])
+    tcs = len([x for x in program.src[0].arg.applied_opts if x.op is OptOps.TC])
     assert wmmas > 0, "tensor core not triggered"
     assert tcs == 1, "tensor core opt not included"
   else:
     try:
-      program = get_program(replace_opts(realized_ast, opts_to_apply), Device[Device.DEFAULT].renderer)
+      program = to_program(replace_opts(realized_ast, opts_to_apply), Device[Device.DEFAULT].renderer)
       assert False, "OptOps.TC triggered, expected KernelOptError"
     except KernelOptError: pass
 
@@ -47,9 +47,10 @@ def helper_tc_allclose(N:int, M:int, K:int, dtype_in:DType, dtype_out:DType, axi
   if dtype_in == dtypes.bfloat16: r = r.float()
   realized_ast, bufs = helper_realized_ast(r)
   opts = [Opt(op=OptOps.TC, axis=axis, arg=(tc_select, tc_opt, use_tensor_cores))]
-  prg = CompiledRunner(replace(get_program(replace_opts(realized_ast, opts), Device[Device.DEFAULT].renderer), device=Device.DEFAULT))
-  if use_tensor_cores == 1: assert len([uop for uop in prg.p.uops if uop.op is Ops.WMMA]) > 0, "wmma not triggered"
-  assert len([x for x in prg.p.uops[-1].arg.applied_opts if x.op is OptOps.TC]) == 1, "tensor core opt not included"
+  pu = to_program(replace_opts(realized_ast, opts), Device[Device.DEFAULT].renderer)
+  if use_tensor_cores == 1: assert len([uop for uop in pu.src[2].src if uop.op is Ops.WMMA]) > 0, "wmma not triggered"
+  assert len([x for x in pu.src[0].arg.applied_opts if x.op is OptOps.TC]) == 1, "tensor core opt not included"
+  prg = CompiledRunner(pu, Device.DEFAULT)
   prg.exec(bufs)
   if dtype_in == dtypes.half: tc_atol, tc_rtol = 1e-2, 1e-3
   elif dtype_in == dtypes.bfloat16: tc_atol, tc_rtol = (1e-1, 2e-2) if dtype_out == dtypes.bfloat16 else (1e-2, 1e-2)
@@ -76,15 +77,16 @@ class TestTensorCores(unittest.TestCase):
       n, m, k = tc.dims[0], tc.dims[1], 2 if AMX else tc.dims[2]
       a, b = Tensor.rand(m, k, dtype=tc.dtype_in), Tensor.rand(k, n, dtype=tc.dtype_in)
       r = a.matmul(b, dtype=tc.dtype_out)
-      prg = get_program(replace_opts(r.schedule()[-1].ast, [Opt(op=OptOps.TC, axis=0, arg=(-1, 2, 1))]), Device[Device.DEFAULT].renderer)
+      prg = to_program(replace_opts(r.schedule_linear().src[-1].src[0],
+                        [Opt(op=OptOps.TC, axis=0, arg=(-1, 2, 1))]), Device[Device.DEFAULT].renderer)
       if Device.DEFAULT == "CPU" and DEV.renderer == "LLVM":
-        assert "0x201000" in prg.src
+        assert "0x201000" in prg.src[3].arg
       elif Device.DEFAULT == "AMD" and DEV.renderer == "LLVM":
-        assert "@llvm.amdgcn.wmma" in prg.src
+        assert "@llvm.amdgcn.wmma" in prg.src[3].arg
       elif Device[Device.DEFAULT].renderer.suffix == "PTX":
-        assert "mma.sync.aligned" in prg.src
+        assert "mma.sync.aligned" in prg.src[3].arg
       else:
-        assert "__WMMA_" in prg.src
+        assert "__WMMA_" in prg.src[3].arg
 
   @Context(ALLOW_TF32=1)
   @unittest.skipIf((Device.DEFAULT == "AMD") or (Device.DEFAULT == "PYTHON" and Device.default.renderer.target.device == "AMD"), "broken for AMD")
@@ -143,11 +145,11 @@ class TestTensorCores(unittest.TestCase):
         c = a.conv2d(b, padding=1, dtype=tc.dtype_out)
         realized_ast, real_bufs = helper_realized_ast(c)
 
-        program = get_program(replace_opts(realized_ast, [Opt(OptOps.TC, axis, (-1, 2, 1))]), Device[Device.DEFAULT].renderer)
-        assert len([uop for uop in program.uops if uop.op is Ops.WMMA]) > 0, "tensor core not triggered"
-        assert len([x for x in program.applied_opts if x.op is OptOps.TC]) == 1, "tensor core opt not included"
+        program = to_program(replace_opts(realized_ast, [Opt(OptOps.TC, axis, (-1, 2, 1))]), Device[Device.DEFAULT].renderer)
+        assert len([uop for uop in tuple(program.src[2].src) if uop.op is Ops.WMMA]) > 0, "tensor core not triggered"
+        assert len([x for x in program.src[0].arg.applied_opts if x.op is OptOps.TC]) == 1, "tensor core opt not included"
 
-        prg = CompiledRunner(program)
+        prg = CompiledRunner(program, Device.DEFAULT)
         # TODO: support this even if numpy doesn't
         if _to_np_dtype(real_bufs[0].dtype) is None: continue
         real_bufs[0].copyin(np.zeros((real_bufs[0].size, ), dtype=_to_np_dtype(real_bufs[0].dtype)).data) # Zero to check that all values are filled
@@ -167,7 +169,7 @@ class TestTensorCores(unittest.TestCase):
     r = x.matmul(y, dtype=tc.dtype_out)
     opts = [Opt(OptOps.UNROLL, 0, 4)]
     ast = helper_linearizer_opt(r, [opts], apply_tc=True, atol=3e-2, rtol=1e-3)
-    for u in get_program(replace_opts(ast, opts), Device[Device.DEFAULT].renderer).uops:
+    for u in tuple(to_program(replace_opts(ast, opts), Device[Device.DEFAULT].renderer).src[2].src):
       if u.op is Ops.WMMA:
         assert u.src[-1].src[0].op != Ops.STORE
 
@@ -181,7 +183,7 @@ class TestTensorCores(unittest.TestCase):
     r = x.matmul(y, dtype=tc.dtype_out)
     opts = [Opt(OptOps.UNROLL, 0, 4)]
     ast = helper_linearizer_opt(r, [opts], apply_tc=True, atol=3e-2, rtol=1e-3)
-    for u in get_program(replace_opts(ast, opts), Device[Device.DEFAULT].renderer).uops:
+    for u in tuple(to_program(replace_opts(ast, opts), Device[Device.DEFAULT].renderer).src[2].src):
       if u.op is Ops.WMMA:
         #assert u.src[-1].dtype == dtypes.float.vec(prod(tc.thread_local_sizes[2]))
         assert u.src[-1].src[0].op != Ops.STORE
@@ -197,7 +199,7 @@ class TestTensorCores(unittest.TestCase):
     r = x.matmul(y, dtype=tc.dtype_out).relu()
     opts = [Opt(OptOps.UNROLL, 0, 4)]
     ast = helper_linearizer_opt(r, [opts], apply_tc=True, atol=3e-2, rtol=1e-3)
-    for u in get_program(replace_opts(ast, opts), Device[Device.DEFAULT].renderer).uops:
+    for u in tuple(to_program(replace_opts(ast, opts), Device[Device.DEFAULT].renderer).src[2].src):
       if u.op is Ops.WMMA:
         #assert u.src[-1].dtype == dtypes.float.vec(prod(tc.thread_local_sizes[2]))
         assert u.src[-1].src[0].op != Ops.STORE

test/testextra/test_export_model.py

@@ -2,6 +2,8 @@ import unittest
 from extra.export_model import export_model, EXPORT_SUPPORTED_DEVICE
 from tinygrad.tensor import Tensor
 from tinygrad.device import Device
+from tinygrad.nn import Linear
+from tinygrad.nn.state import get_state_dict
 from tinygrad import dtypes
 import json
 
@@ -66,5 +68,15 @@ class TextModelExportWebGPU(unittest.TestCase):
       self.assertIn(f"const resultBuffer{i} = new {expected_arr_prefix}Array(gpuReadBuffer{i}.size/{dt.itemsize});", prg)
       self.assertIn(f"resultBuffer{i}.set(new {expected_arr_prefix}Array(gpuReadBuffer{i}.getMappedRange()));", prg)
 
+  def test_weights_bound_to_safetensor(self):
+    # regression test: every weight ended up as createEmptyBuf (zero-init) instead of createWeightBuf
+    class MyModel:
+      def __init__(self): self.fc1, self.fc2 = Linear(4, 8), Linear(8, 2)
+      def forward(self, x): return self.fc2(self.fc1(x).relu())
+    model = MyModel()
+    for t in get_state_dict(model).values(): t.realize()
+    prg, _, _, _ = export_model(model, "webgpu", Tensor.randn(1, 4))
+    self.assertEqual(prg.count("createWeightBuf("), len(get_state_dict(model)))
+
 if __name__ == '__main__':
   unittest.main()