work

* local optimize as rewrite

* better

* x

* slighly rename

* fix

* ugh

* remove

* x

* remove

* not weak

.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

.github/workflows/test.yml

@@ -676,8 +676,8 @@ jobs:
         env:
           AMD: 0
         run: |
-          PYTHONPATH=. DEV=NULL::gfx1100 python extra/mmapeak/mmapeak.py
-          PYTHONPATH=. DEV=NULL::gfx1201 python3 -m pytest -n=auto test/testextra/test_tk.py test/backend/test_asm_gemm.py
+          PYTHONPATH=. DEV=NULL:HIP:gfx1100 python extra/mmapeak/mmapeak.py
+          PYTHONPATH=. DEV=NULL:HIP:gfx950 python3 -m pytest -n=auto test/testextra/test_tk.py test/backend/test_asm_gemm.py
       - name: Run matmul on MOCKKFD
         run: |
           PYTHONPATH="." DEV=MOCKKFD+AMD N=256 python3 extra/gemm/amd_asm_matmul.py
@@ -773,7 +773,7 @@ jobs:
           cuda: 'true'
           ocelot: 'true'
       - name: Set env
-        run: printf "${{ matrix.backend == 'ptx' && 'DEV=MOCK+CUDA:PTX' || matrix.backend == 'nv' && 'DEV=MOCKNVK+NV\nSKIP_SLOW_TEST=1' }}" >> $GITHUB_ENV
+        run: printf "${{ matrix.backend == 'ptx' && 'DEV=MOCK+CUDA:PTX' || matrix.backend == 'nv' && 'DEV=MOCK+NV\nSKIP_SLOW_TEST=1' }}" >> $GITHUB_ENV
       - name: Check Device.DEFAULT and print some source
         run: |
           python3 -c "from tinygrad import Device; assert Device.DEFAULT in ['CUDA','NV'], Device.DEFAULT"
@@ -870,7 +870,7 @@ jobs:
         python -m pytest -n=auto test/device/test_hcq.py test/test_tiny.py test/device/test_amd_llvm.py --durations=20
     - name: Run pytest (ptx)
       env:
-        DEV: "MOCKNVK+NV:PTX"
+        DEV: "MOCK+NV:PTX"
         FORWARD_ONLY: 1
         # TODO: failing due to library loading error
         CAPTURE_PROCESS_REPLAY: 0

.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,15 +17,13 @@ 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_schedule
+::: tinygrad.engine.realize.run_linear
 
 There's a ton of complexity hidden behind this, see the `codegen/` directory.
 
@@ -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/env_vars.md

@@ -57,6 +57,8 @@ AMD:LLVM      | use the AMD device with the LLVM renderer
 NV:CUDA:sm_70 | use the NV device with the CUDA renderer targetting sm_70
 AMD::gfx950   | use the AMD device targetting gfx950
 USB+AMD       | use the AMD device over the USB interface
+CPU:LLVM      | use the CPU device with the LLVM renderer
+CPU:LLVM:x86_64,znver2,avx2,-avx512f | use the CPU device with the LLVM renderer, with [additional arch flags](runtime.md#cpu-arch)
 
 ### Debug breakdown
 

docs/runtime.md

@@ -10,7 +10,7 @@ tinygrad supports various runtimes, enabling your code to scale across a wide ra
 | [METAL](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_metal.py) | Utilizes Metal for acceleration on Apple devices | - | M1+ Macs; Metal 3.0+ for `bfloat` support |
 | [CUDA](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_cuda.py) | Utilizes CUDA for acceleration on NVIDIA GPUs | nvrtc (default)<br> PTX (`DEV=CUDA:PTX`) | NVIDIA GPU with CUDA support |
 | [CL](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_cl.py) | Accelerates computations using OpenCL on GPUs | - | OpenCL 2.0 compatible device |
-| [CPU](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_cpu.py) | Runs on CPU using the clang or llvm compiler | Clang JIT (default)<br>LLVM IR (`DEV=CPU:LLVM`) | `clang` compiler in system `PATH` |
+| [CPU](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_cpu.py) | Runs on CPU using the clang or llvm compiler | Clang JIT (default)<br>LLVM IR (`DEV=CPU:LLVM`) | `clang` compiler in system `PATH`<br>You can specify additional arch parameters via [the `DEV` variable](env_vars.md#dev-variable). See [CPU arch](#cpu-arch) for details. |
 | [WEBGPU](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_webgpu.py) | Runs on GPU using the Dawn WebGPU engine (used in Google Chrome) | - | Dawn library installed and discoverable. Binaries: [pydawn v0.3.0](https://github.com/wpmed92/pydawn/releases/tag/v0.3.0) |
 
 
@@ -79,3 +79,13 @@ NV backend supports several interfaces for communicating with devices:
 
 * `NVK`: uses the nvidia driver
 * `PCI`: uses the [NV driver](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/support/nv/nvdev.py)
+
+## CPU Arch
+The CPU renderers may be additionally configured using the arch component of [the `DEV` environment variable](env_vars.md#dev-variable).
+CPU arch should be specified as a comma-separated list of parameters, and must contain at least two values: the architecture family (ie. x86_64, arm64, or riscv64) and the cpu type (as accepted by `clang`'s `-march`).
+If native is specified as the cpu type, tinygrad (or delegate compiler) will query the host cpu type. Additional comma-separated values may be specified as follows:
+
+* `AMX`: emit Apple silicon AMX instructions
+
+All other additional values are interpreted as cpu feature flags. When a value is preceded by a `-` character, the corresponding feature flag will be disabled, otherwise the flag will be enabled.
+Note that enabled feature flags should not be preceded by a `+`.

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

@@ -113,7 +113,7 @@ class VLIWRenderer(Renderer):
         case Ops.GEP:
           # a GEP is just an alias to a special register in the vector
           r[u] = r[u.src[0]] + u.arg[0]
-        case Ops.VECTORIZE:
+        case Ops.STACK:
           if all(s == u.src[0] for s in u.src):
             # if all sources are the same, we can broadcast
             inst.append({"valu": [("vbroadcast", r[u], r[u.src[0]])]})
@@ -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,11 +128,8 @@ class FlatTransformer:
     bsz, seqlen, _ = x.shape
     new_amaxs, saves = [], []
 
-    x, rrms = rmsnorm(x, self.norm_eps)
-    saves.extend([x, rrms])
-    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)
@@ -148,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
@@ -167,22 +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])
-    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)
@@ -213,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, Ops.BEAM}
+  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/tools/gpuburn.py

@@ -1,7 +1,6 @@
 from tinygrad import Tensor, Device, TinyJit, dtypes
-from tinygrad.helpers import getenv
 
-GPUS = getenv("GPUS", 4) # TODO: expose a way in tinygrad to access this
+GPUS = Device[Device.DEFAULT].count()
 N = 6144
 
 @TinyJit

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

@@ -6,13 +6,13 @@ 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
 
 EXPORT_SUPPORTED_DEVICE = ["WEBGPU", "CPU", "CUDA", "CL"]
 
-_KERNEL_ASTS = {Ops.SINK, Ops.PROGRAM, Ops.BEAM}
+_KERNEL_ASTS = {Ops.SINK, Ops.PROGRAM}
 def iter_kernel_calls(linear:UOp):
   """Yield kernel CALLs from a LINEAR UOp. Toposort descends naturally into CUSTOM_FUNCTION graph batches; gate stops at kernel ASTs."""
   return (u for u in 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)
@@ -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/cdna_asm_gemm.py

@@ -2745,14 +2745,14 @@ def asm_gemm(a:Tensor, b:Tensor, x_scale:Tensor|None=None, w_scale:Tensor|None=N
 
   if is_multi:
     if n_sharded:
-      out = Tensor(Tensor.invalid(batch, M, N//len(a.device), dtype=out_dtype, device=a.device).uop.multi(2), device=a.device)
+      out = Tensor(Tensor.invalids(batch, M, N//len(a.device), dtype=out_dtype, device=a.device).uop.multi(2), device=a.device)
     elif m_sharded:
-      out = Tensor(Tensor.invalid(batch, M, N, dtype=out_dtype, device=a.device).uop.multi(1), device=a.device)
+      out = Tensor(Tensor.invalids(batch, M, N, dtype=out_dtype, device=a.device).uop.multi(1), device=a.device)
     else:
-      out = Tensor(Tensor.invalid(batch//len(a.device) if a.uop.axis==0 else batch, M, N, dtype=out_dtype, device=a.device).uop.multi(0),
+      out = Tensor(Tensor.invalids(batch//len(a.device) if a.uop.axis==0 else batch, M, N, dtype=out_dtype, device=a.device).uop.multi(0),
                    device=a.device)
   else:
-    out = Tensor.invalid(batch, M, N, dtype=out_dtype, device=a.device)
+    out = Tensor.invalids(batch, M, N, dtype=out_dtype, device=a.device)
 
   renderer = Device[dname:=(a.device[0] if is_multi else a.device)].renderer
   dname, arch = dname.split(":")[0], renderer.target.arch

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,8 +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)
-  sched = C.schedule()
-  si = sched[-1]
+  from tinygrad.uop.ops import Ops
+  linear, var_vals = C.linear_with_vars()
+  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
@@ -85,22 +89,25 @@ 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}")
 
   # check correctness
   if getenv("VERIFY"):
-    from tinygrad.engine.realize import run_schedule
+    from tinygrad.engine.realize import run_linear
     triton_buf = np.frombuffer(si.bufs[0].as_memoryview(), np.float16).reshape(M,N)
     print(triton_buf)
-    run_schedule(sched)
+    run_linear(linear, var_vals)
     tinygrad_buf = np.frombuffer(si.bufs[0].as_memoryview(), np.float16).reshape(M,N)
     print(tinygrad_buf)
     np.testing.assert_allclose(triton_buf, tinygrad_buf)

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

@@ -1,30 +1,19 @@
+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
-
-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
+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")
-  # read 2*N bf16 (xw13) + N bf16 (grad_x2) + 1 scalar; write 2*N bf16 (grad_xw13)
   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)
+  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)))
 
@@ -33,33 +22,34 @@ def _custom_fused_cast_amax_w13(fp8_out:UOp, amax_buf:UOp, xw13:UOp, amax_state:
   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")
-  # read 2*N bf16 + 1 scalar, write N fp8 + NUM_WG bf16
   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)
+  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):
-  # kernel.src[1:] is (fp8_out, amax_buf, xw13, amax_state); only xw13 needs a grad
+  # 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.invalid(*_shard_shape(xw13.shape, axis, ndev), dtype=dtypes.bfloat16, device=device).uop.multi(axis), device=device)
+    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.invalid(*xw13.shape, dtype=dtypes.bfloat16, device=device)
+    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)
+  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).
+  # 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}"
@@ -67,19 +57,17 @@ def fused_quantize_fp8_w13(xw13:Tensor, amax_state:Tensor, fp8_dtype) -> tuple[T
   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.invalid(*_shard_shape((MBS, SEQ, HIDDEN), axis, ndev), dtype=fp8_dtype, device=xw13.device).uop.multi(axis), device=xw13.device)
-    amax_buf = Tensor(Tensor.invalid(NUM_WG, dtype=dtypes.bfloat16, device=xw13.device).uop.multi(0), device=xw13.device)
+    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.invalid(MBS, SEQ, HIDDEN, dtype=fp8_dtype, device=xw13.device)
-    amax_buf = Tensor.invalid(NUM_WG, dtype=dtypes.bfloat16, device=xw13.device)
+    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)
-  # per-device scalar amax (no cross-device allreduce, matches _local_abs_max semantics)
-  if isinstance(amax_buf.device, tuple):
-    from examples.mlperf.models.flat_llama import _local_abs_max
-    new_amax = _local_abs_max(amax_buf).detach()
-  else: new_amax = amax_buf.max().detach()
-  inv_scale = (FP8_MAX / (amax_state + 1e-8)).float().reciprocal()
-  return fp8_out, inv_scale, new_amax
+  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/fused_mul_quantize_fp8/fused_mul_quantize_fp8.cpp

@@ -0,0 +1,71 @@
+#include <hip/hip_runtime.h>
+#include <hip/hip_bf16.h>
+#include <hip/hip_fp8.h>
+
+#ifndef N_ELEMS
+#define N_ELEMS 67108864
+#endif
+#ifndef HIDDEN
+#define HIDDEN 4096
+#endif
+#ifndef NUM_WG
+#define NUM_WG 1024
+#endif
+#ifndef THREADS_PER_WG
+#define THREADS_PER_WG 256
+#endif
+
+constexpr int VEC = 8;
+constexpr float FP8_MAX = 448.0f;
+
+static_assert(N_ELEMS % VEC == 0, "N_ELEMS must be divisible by VEC");
+static_assert(HIDDEN % VEC == 0, "HIDDEN must be divisible by VEC");
+
+extern "C" __global__ __launch_bounds__(THREADS_PER_WG) void
+fused_mul_quantize_fp8(
+    __hip_fp8_storage_t*  __restrict__ fp8_out,         // fp8, N_ELEMS
+    __hip_bfloat16*       __restrict__ amax_buf,        // bf16, NUM_WG
+    const __hip_bfloat16* __restrict__ x,               // bf16, N_ELEMS
+    const __hip_bfloat16* __restrict__ weight,          // bf16, HIDDEN (per-hidden scale)
+    const __hip_bfloat16* __restrict__ amax_state)      // bf16 scalar
+{
+  __shared__ float sdata[THREADS_PER_WG];
+
+  const int tid = threadIdx.x;
+  const int wg  = blockIdx.x;
+  const int gid = wg * THREADS_PER_WG + tid;
+  const int stride_elems = NUM_WG * THREADS_PER_WG * VEC;
+
+  const float scale = FP8_MAX / (static_cast<float>(*amax_state) + 1e-8f);
+  float local_max = 0.0f;
+
+  for (int base = gid * VEC; base < N_ELEMS; base += stride_elems) {
+    const int h = base % HIDDEN;   // 0..HIDDEN-VEC, 8-aligned (since base is 8-aligned and HIDDEN divides VEC)
+    float4 x_raw = *reinterpret_cast<const float4*>(&x[base]);
+    float4 w_raw = *reinterpret_cast<const float4*>(&weight[h]);
+
+    const __hip_bfloat16 *xi = reinterpret_cast<const __hip_bfloat16*>(&x_raw);
+    const __hip_bfloat16 *wi = reinterpret_cast<const __hip_bfloat16*>(&w_raw);
+
+    __hip_fp8_storage_t out[VEC];
+    #pragma unroll
+    for (int i = 0; i < VEC; i++) {
+      const float val = static_cast<float>(xi[i]) * static_cast<float>(wi[i]);
+      local_max = fmaxf(local_max, fabsf(val));
+      const float scaled = fmaxf(-FP8_MAX, fminf(FP8_MAX, val * scale));
+      out[i] = __hip_cvt_float_to_fp8(scaled, __HIP_SATFINITE, __HIP_E4M3);
+    }
+
+    *reinterpret_cast<uint64_t*>(&fp8_out[base]) = *reinterpret_cast<uint64_t*>(out);
+  }
+
+  // LDS tree-reduce per-WG amax
+  sdata[tid] = local_max;
+  __syncthreads();
+  for (int s = THREADS_PER_WG / 2; s > 0; s >>= 1) {
+    if (tid < s) sdata[tid] = fmaxf(sdata[tid], sdata[tid + s]);
+    __syncthreads();
+  }
+
+  if (tid == 0) amax_buf[wg] = static_cast<__hip_bfloat16>(sdata[0]);
+}

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

@@ -10,4 +10,4 @@ def extract_ast(*args) -> None:
   return None
 
 if __name__ == "__main__":
-  _pmap({"get_program":extract_ast})
+  _pmap({"do_to_program":extract_ast})

extra/optimization/test_beam_search.py

@@ -84,8 +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
-    if ast.op is Ops.BEAM: ast = ast.src[0]
+    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)])
@@ -95,8 +94,7 @@ class TestBeamSearch(unittest.TestCase):
 
   def test_max_up(self):
     a = Tensor.rand(16, 16)
-    ast = a.schedule()[-1].ast
-    if ast.op is Ops.BEAM: ast = ast.src[0]
+    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)

extra/sqtt/install_rocprof_decoder.py

@@ -1,20 +1,23 @@
 #!/usr/bin/env python3
-import os, shutil
+import os, platform, shutil, subprocess
 from pathlib import Path
 from tinygrad.helpers import fetch, OSX
 
+VERSION = "0.1.6"
 DEST = Path("/usr/local/lib")
 DEST.mkdir(exist_ok=True)
 
 if __name__ == "__main__":
   if OSX:
-    fp = fetch("https://github.com/ROCm/rocprof-trace-decoder/releases/download/0.1.4/rocprof-trace-decoder-macos-arm64-0.1.4-Darwin.sh")
-    lib = fp.parent/"rocprof-trace-decoder-macos-arm64-0.1.4-Darwin"/"lib"/"librocprof-trace-decoder.dylib"
-    os.chmod(fp, 0o755)
-    os.system(f"sudo {fp} --prefix={fp.parent} --include-subdir")
-    shutil.copy2(lib, DEST)
+    arch = "arm64" if platform.machine() == "arm64" else "x86_64"
+    dmg = fetch(f"https://github.com/ROCm/rocprof-trace-decoder/releases/download/{VERSION}/rocprof-trace-decoder-macos-{arch}-{VERSION}-Darwin.dmg")
+    mnt = Path(subprocess.check_output(["hdiutil", "attach", "-nobrowse", "-readonly", "-mountrandom", "/tmp", str(dmg)],
+                                       text=True).split("\t")[-1].strip())
+    try: shutil.copy2(next(mnt.rglob("librocprof-trace-decoder.dylib")), DEST)
+    finally: subprocess.run(["hdiutil", "detach", str(mnt)], check=True)
+    lib = DEST/"librocprof-trace-decoder.dylib"
   else:
     lib = DEST/"librocprof-trace-decoder.so"
-    os.system("sudo curl -L https://github.com/ROCm/rocprof-trace-decoder/raw/43bf0fef74a83c3c25badfc5a09c0bd39ed8c6f9/releases/linux_glibc_2_28_x86_64/librocprof-trace-decoder.so -o"+str(lib))
+    os.system(f"sudo curl -L https://github.com/ROCm/rocprof-trace-decoder/raw/{VERSION}/releases/linux_glibc_2_28_x86_64/librocprof-trace-decoder.so -o {lib}")
     os.system("sudo ldconfig")
-  print(f"Installed {lib.name} to", DEST)
+  print(f"Installed {lib.name} ({VERSION}) to", DEST)

extra/thunder/amd/fa.py

@@ -10,11 +10,11 @@ from tinygrad.uop.ops import UOp, Ops, KernelInfo
 
 def _sharded_empty(shape:Tensor, ref:Tensor, axis:int|None, dtype:DTypeLike|None=None) -> Tensor:
   dtype = dtype or ref.dtype
-  if not isinstance(ref.device, tuple): return Tensor.invalid(*shape, dtype=dtype, device=ref.device)
+  if not isinstance(ref.device, tuple): return Tensor.invalids(*shape, dtype=dtype, device=ref.device)
   shard_axis = ref.uop.axis if axis is None else axis
   shape = tuple(s // len(ref.device) if i == shard_axis else s for i, s in enumerate(shape))
   axis = ref.uop.axis if axis is None else axis
-  return Tensor(Tensor.invalid(*shape, dtype=dtype, device=ref.device).uop.multi(axis), dtype=dtype, device=ref.device)
+  return Tensor(Tensor.invalids(*shape, dtype=dtype, device=ref.device).uop.multi(axis), dtype=dtype, device=ref.device)
 
 def _sharded_empty_like(ref:Tensor, axis:int|None=None) -> Tensor:
   return _sharded_empty(ref.shape, ref, axis)

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

@@ -1,7 +1,7 @@
 #!/usr/bin/env python3
 """Tests for SQTT encoder: verifies the emulator produces correct SQTT traces for known kernels.
 
-Run with: DEV=MOCKKFD+AMD python -m pytest test/amd/test_sqtt_encoder.py -v
+Run with: DEV=MOCK+AMD python -m pytest test/amd/test_sqtt_encoder.py -v
 """
 import ctypes, unittest
 from tinygrad.helpers import Context

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,22 +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_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)
@@ -40,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
 
@@ -54,9 +49,9 @@ class TestIndexing(unittest.TestCase):
     with Context(NOOPT=1):
       GlobalCounters.reset()
       out = ((Tensor.arange(1,16385)-1)*needle).sum()
-      sched = out.schedule()
-      self.assertEqual(len(sched), 1)
-      run_schedule(sched)
+      linear, var_vals = out.linear_with_vars()
+      self.assertEqual(len(linear.src), 1)
+      run_linear(linear, var_vals)
     self.assertEqual(out.item(), 1337)
 
   def test_manual_index(self):
@@ -71,9 +66,9 @@ class TestIndexing(unittest.TestCase):
       reshape_dataset = dataset.T.reshape(1, DDIM, DSET, 1).expand(4, DDIM, DSET, 1)
       full = (rng==idxs).where(reshape_dataset, Tensor.zeros(4, DDIM, DSET, 1))
       X = full.sum(axis=(2,3))
-      sched = X.schedule()
-      self.assertEqual(len(sched), 1)
-      run_schedule(sched)
+      linear, var_vals = X.linear_with_vars()
+      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())
 
@@ -97,9 +92,9 @@ class TestIndexing(unittest.TestCase):
       GlobalCounters.reset()
       X = dataset[idxs]
       assert X.shape == (4,DDIM)
-      sched = X.schedule()
-      self.assertEqual(len(sched), 1)
-      run_schedule(sched)
+      linear, var_vals = X.linear_with_vars()
+      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())
 
@@ -112,9 +107,9 @@ class TestIndexing(unittest.TestCase):
       GlobalCounters.reset()
       X = dataset[idxs]
       assert X.shape == (4,DDIM)
-      sched = X.schedule()
-      self.assertEqual(len(sched), 1)
-      run_schedule(sched)
+      linear, var_vals = X.linear_with_vars()
+      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())
   @unittest.skip("not ready")
@@ -234,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}")
@@ -299,8 +299,8 @@ class TestCustomKernel(unittest.TestCase):
     from tinygrad import function
     @function(precompile=True)
     def run(x:Tensor, w:Tensor) -> Tensor:
-      out = Tensor.invalid(*x.shape, dtype=x.dtype)
-      tmp = Tensor.invalid(*x.shape, dtype=x.dtype)
+      out = Tensor.invalids(*x.shape, dtype=x.dtype)
+      tmp = Tensor.invalids(*x.shape, dtype=x.dtype)
       out, tmp = Tensor.custom_kernel(out, tmp, x, w, fxn=custom_add_with_tmp)[:2]
       return out+tmp
 

test/backend/test_device.py

@@ -0,0 +1,9 @@
+import unittest
+from tinygrad import Device
+
+class TestDeviceCount(unittest.TestCase):
+  def test_count(self):
+    self.assertGreaterEqual(Device[Device.DEFAULT].count(), 1)
+
+if __name__ == "__main__":
+  unittest.main()

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.schedule import linear_to_schedule
+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,12 +44,12 @@ 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, [])
 
 def run_schedule(calls:list[UOp]):
-  for ei in linear_to_schedule(UOp(Ops.LINEAR, src=tuple(calls))): ei.lower().run({})
+  run_linear(UOp(Ops.LINEAR, src=tuple(calls)))
 
 def zero_bufs(bufs):
   for b in bufs:

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
@@ -706,7 +720,7 @@ class TestJitGraphSplit(unittest.TestCase):
         inner_cnt = len(ast.src[0].src)
         assert inner_cnt == expected["cnt"], f"Expected {expected['cnt']} operations in graph, got {inner_cnt}"
       elif expected["type"] == "comp":
-        assert ast.op in (Ops.SINK, Ops.PROGRAM, Ops.BEAM), f"Expected kernel, got {ast.op}"
+        assert ast.op in (Ops.SINK, Ops.PROGRAM), f"Expected kernel, got {ast.op}"
       elif expected["type"] in ("copy", "xfer"):
         assert ast.op is Ops.COPY, f"Expected COPY, got {ast.op}"
 

test/backend/test_linearizer.py

@@ -1,12 +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_schedule, CompiledRunner, get_program
+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
@@ -25,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)
@@ -45,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
@@ -53,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")
@@ -63,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
@@ -84,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
 
@@ -92,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
 
@@ -100,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]])
@@ -113,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
@@ -123,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
 
@@ -134,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?"
@@ -146,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"
 
@@ -155,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
@@ -168,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")
@@ -179,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)
@@ -193,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"
 
@@ -202,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 = (
@@ -238,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)
@@ -258,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]
@@ -267,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"
 
@@ -286,18 +290,17 @@ class TestLinearizer(unittest.TestCase):
     a = Tensor.ones(4, 4).contiguous().realize()
     b = a.shrink(((1, 2), None)).pad(((1, 2), None))
     a.assign(b.where(2, a))
-    sched = a.schedule()
-    assert len(sched) == 1
-    sched_copy = sched[:]
-    run_schedule(sched)
+    linear, var_vals = a.linear_with_vars()
+    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)]
@@ -329,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:
@@ -340,8 +343,8 @@ 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]
-    assert store_val.dtype == dtypes.float.vec(4) and store_val.op is not Ops.VECTORIZE
+    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")
   @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_shared, "test requires shared")
@@ -353,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]
@@ -374,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
@@ -388,15 +391,18 @@ class TestLinearizer(unittest.TestCase):
 
 def helper_realized_ast(r:Tensor|list[Tensor]) -> tuple[UOp, list[Buffer]]:
   if isinstance(r, Tensor): r = [r]
-  s = Tensor.schedule(*r)
-  run_schedule(s[:-1])  # 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
+  linear, var_vals = Tensor.linear_with_vars(*r)
+  run_linear(UOp(Ops.LINEAR, src=linear.src[:-1]), var_vals)  # run all kernels except the last one
+  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"
@@ -423,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,7 +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_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
@@ -127,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):
@@ -192,11 +189,11 @@ class TestMultiTensor(unittest.TestCase):
     # only shrink on the device that owns the shard, this is enabled by the mselect simplifier
     for i in range(2):
       xt = X[i*2:i*2+2].contiguous()
-      sched = xt.schedule()
-      #kernels = [s for s in sched if s.ast.op is Ops.SINK]
+      linear, var_vals = xt.linear_with_vars()
+      #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])
-      run_schedule(sched)
+      #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])
 
   @given(strat.sampled_from((devices_2, devices_3)),
@@ -275,6 +272,14 @@ class TestMultiTensor(unittest.TestCase):
       out = f(tt)
       assert out.item() == 1+2+3+4
 
+  def test_multitensor_jit_input_reduce_shard_axis(self):
+    @TinyJit
+    def f(x): return x.sum(0).realize()
+    for _ in range(5):
+      tt = Tensor.ones(2, 64).contiguous().realize().shard((d1,d2), 0).realize()
+      out = f(tt)
+      np.testing.assert_allclose(out.numpy(), np.full(64, 2.0))
+
   def test_multitensor_inside_jit(self):
     @TinyJit
     def f(x): return (x.shard((d1,d2), 0)+1).contiguous().sum()
@@ -546,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:
@@ -631,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):
@@ -676,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
@@ -687,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
@@ -784,9 +801,9 @@ class TestMultiTensor(unittest.TestCase):
   def test_full_like_shrink_on_shard_axis(self):
     t = Tensor.ones(16, 16, dtype=dtypes.int).shard(devices_2, axis=0)
     out = Tensor.full_like(t, 2)[:, :8]
-    sched = out.schedule()
-    self.assertEqual(len(sched), 0)
-    run_schedule(sched)
+    linear, var_vals = out.linear_with_vars()
+    self.assertEqual(len(linear.src), 0)
+    run_linear(linear, var_vals)
     self.assertEqual(out.tolist(), [[2]*8]*16)
 
   def test_dropout_on_shard(self):
@@ -828,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])
 
@@ -839,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()
@@ -849,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
@@ -870,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]))
@@ -1096,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
@@ -1135,13 +1152,13 @@ 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()
-    sched = b_multi.schedule()
-    compiled = [si for si in sched if isinstance(si.prg, CompiledRunner)]
-    self.assertEqual(len(compiled), 0, f"expected zero compiled kernels, got {len(compiled)}")
-    run_schedule(sched)
+    linear, var_vals = b_multi.linear_with_vars()
+    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())
 
   @unittest.skip("flaky on LLVM")
@@ -1168,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()
-    sched = a[5].contiguous().schedule()
-    compiled = [si for si in sched if isinstance(si.prg, CompiledRunner)]
-    self.assertEqual(len(compiled), 0)
-    run_schedule(sched)
-    np.testing.assert_equal(a[5].contiguous().numpy(), ref[5].numpy())
+    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(out.numpy(), ref[5].numpy())
 
 @unittest.skipIf(not_support_multi_device(), "need multi")
 class TestMultiFromUnrenderable(unittest.TestCase):

test/backend/test_nn.py

@@ -8,7 +8,7 @@ from tinygrad.helpers import GlobalCounters, Context
 from tinygrad.nn import Conv1d, ConvTranspose1d, Conv2d, ConvTranspose2d, Linear, Embedding
 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_schedule
+from tinygrad.engine.realize import run_linear
 from test.helpers import not_support_multi_device, needs_second_gpu, slow
 
 @slow
@@ -431,17 +431,19 @@ class TestNN(unittest.TestCase):
     a = Tensor([[1, 5, 9, 11],
                 [12, 19, 8, 1]])
     result = layer(a)
-    schedule = result.schedule()
-    self.assertEqual(len([item for item in schedule if item.ast.op is Ops.SINK]), kcount, "first run realizes weight and embedding")
-    run_schedule(schedule)
+    linear, var_vals = result.linear_with_vars()
+    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)
 
     b = Tensor([[1, 2, 3],
                 [4, 5, 6],
                 [7, 8, 9]])
     result = layer(b)
-    schedule = result.schedule()
-    self.assertEqual(1, len([item for item in schedule if item.ast.op is Ops.SINK]), "second run realizes embedding only")
-    run_schedule(schedule)
+    linear, var_vals = result.linear_with_vars()
+    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")
     self.assertLessEqual(GlobalCounters.global_ops, ops)
 

test/backend/test_ops.py

@@ -3298,13 +3298,11 @@ class TestOps(unittest.TestCase):
     data = [1, 2, 4]
     helper_test_op([], lambda: torch.nn.functional.one_hot(torch.tensor(data), 6).type(torch.int32),
                        lambda: Tensor(data).one_hot(6), forward_only=True)
-    helper_test_op([], lambda: torch.nn.functional.one_hot(torch.tensor(data)).type(torch.int32),
-                       lambda: Tensor(data).one_hot(), forward_only=True)
+    # like jax.nn.one_hot, num_classes must be non-negative (torch accepts -1 for auto-inference, we don't)
+    with self.assertRaises(ValueError): Tensor(data).one_hot(-1)
     data = [[[1, 2, 3], [0, 3, 5]], [[1, 2, 3], [0, 3, 5]]]
     helper_test_op([], lambda: torch.nn.functional.one_hot(torch.tensor(data), 8).type(torch.int32),
                        lambda: Tensor(data).one_hot(8), forward_only=True)
-    helper_test_op([], lambda: torch.nn.functional.one_hot(torch.tensor(data)).type(torch.int32),
-                       lambda: Tensor(data).one_hot(), forward_only=True)
 
   def test_masked_fill(self):
     helper_test_op([(32,10)], lambda x: x.masked_fill((x>0.1).detach(), -math.inf))

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_schedule, 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):
@@ -49,9 +48,9 @@ def _test_conv2d(allowed:int, dtype:DType=dtypes.float):
   w = Tensor.uniform(16, CIN, 3, 3, requires_grad=True).realize()
   ret = Tensor.conv2d(img, w).relu().mean().backward()
   dtypes.default_float = old_default_float
-  s = Tensor.schedule(ret, img.grad, w.grad)
-  run_schedule(s.copy())
-  cnt = len([si for si in s if si.ast.op is Ops.SINK])
+  linear, var_vals = Tensor.linear_with_vars(ret, img.grad, w.grad)
+  run_linear(linear, var_vals)
+  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
@@ -72,9 +71,9 @@ class TestSchedule(unittest.TestCase):
   def test_arange_avgpool2d(self, kcount=1):
     x = Tensor.arange(25).reshape(1,1,5,5).cast(dtypes.float32)
     t = x.avg_pool2d(padding=1)
-    sched = t.schedule()
-    self.assertEqual(len(sched), kcount)
-    run_schedule(sched)
+    linear, var_vals = t.linear_with_vars()
+    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()
     np.testing.assert_allclose(t.numpy(), torch_out)
@@ -788,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)
 
@@ -798,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)
@@ -815,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)
@@ -829,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)
@@ -1053,8 +1049,9 @@ class TestSchedule(unittest.TestCase):
     a2 = mop(a)
     expected = (a+a2).tolist()
     a.assign(a+a2)
-    kcount = len(sched:=a.schedule())
-    run_schedule(sched)
+    linear, var_vals = a.linear_with_vars()
+    kcount = len(linear.src)
+    run_linear(linear, var_vals)
     self.assertListEqual(a.tolist(), expected)
     self.assertEqual(kcount, expected_kcount)
   def test_setitem_permuted_sched(self): self.test_setitem_sched(lambda x: x.T, 2)
@@ -1332,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()
@@ -1353,9 +1350,9 @@ class TestCopyFolding(unittest.TestCase):
   def test_copy_to_same_device_sched(self):
     a = Tensor.ones(4).contiguous().realize().uop.buf_uop
     t = Tensor(a.copy_to_device(a.device))
-    sched = t.schedule()
-    assert len([s for s in sched if s.ast.op is Ops.COPY]) == 0
-    run_schedule(sched)
+    linear, var_vals = t.linear_with_vars()
+    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.])
 
@@ -1442,8 +1439,7 @@ class TestFusionOp(unittest.TestCase):
   def test_expand_fuse(self):
     bt = Tensor(np.ones((10, 1)), dtype=dtypes.float32)
     out = (bt*2).expand(10,10).sum(1)
-    sched = out.schedule()
-    run_schedule(sched)
+    run_linear(*out.linear_with_vars())
     outd = out.tolist()
     assert all(x == 20.0 for x in outd)
 

test/backend/test_setitem.py

@@ -293,7 +293,7 @@ class TestSetitem(unittest.TestCase):
     np.testing.assert_allclose(b.numpy(), [0, 2, 4, 6])
 
   def test_setitem_multiple_disjoint_on_invalid(self):
-    z = Tensor.invalid(10, dtype="int").realize()
+    z = Tensor.invalids(10, dtype="int").realize()
     z[2:5] = 2
     z[6:7] = 3
     z.realize()

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

@@ -6,7 +6,7 @@ from tinygrad import Tensor, TinyJit, GlobalCounters, Device
 from tinygrad.helpers import getenv, Context
 from tinygrad.nn.optim import SGD
 from tinygrad.nn.state import get_parameters
-from tinygrad.engine.realize import run_schedule
+from tinygrad.engine.realize import run_linear
 
 from extra.models import resnet
 from examples.mlperf.initializers import Conv2dHeNormal, Linear
@@ -71,11 +71,11 @@ class BenchmarkResnetTrain(unittest.TestCase):
 
       y = x.sequential(layer).contiguous().contiguous_backward()
       y.sum().backward()
-      if getenv("ASSIGN", 1): sched, _ = Tensor.schedule_with_vars(y, x.grad, *optim.schedule_step())
-      else: sched, _ = Tensor.schedule_with_vars(y, x.grad, *[t.grad for t in optim.params])
+      if getenv("ASSIGN", 1): linear, var_vals = Tensor.linear_with_vars(y, x.grad, *optim.schedule_step())
+      else: linear, var_vals = Tensor.linear_with_vars(y, x.grad, *[t.grad for t in optim.params])
 
       for _ in range(JITCNT):
-        run_schedule(list(sched))
+        run_linear(linear, var_vals)
 
     CNT = getenv("CNT", 5)
     best_tm = None

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

@@ -6,8 +6,7 @@ import torch
 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_schedule
-from tinygrad.schedule import linear_to_schedule
+from tinygrad.engine.realize import capturing, run_linear
 from tinygrad.tensor import _to_np_dtype
 
 class CLCache:
@@ -15,9 +14,8 @@ 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):
-    schedule = linear_to_schedule(linear)
-    self.count += len(schedule)
-    run_schedule(schedule, var_vals)
+    self.count += len(linear.src)
+    run_linear(linear, var_vals)
   def __enter__(self):
     if self.preclear:
       gc.collect()

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

@@ -4,7 +4,7 @@ from tinygrad import Tensor, TinyJit, GlobalCounters, Device
 from tinygrad.helpers import getenv, Context
 from tinygrad.nn.optim import LAMB
 from tinygrad.nn.state import get_parameters
-from tinygrad.engine.realize import run_schedule
+from tinygrad.engine.realize import run_linear
 
 from extra.models import bert
 
@@ -49,11 +49,11 @@ class BenchmarkBertTrain(unittest.TestCase):
 
       y = layer(*inputs).contiguous().contiguous_backward()
       y.sum().backward()
-      if getenv("ASSIGN", 1): sched, _ = Tensor.schedule_with_vars(y, *list(inputs), *optim.schedule_step())
-      else: sched, _ = Tensor.schedule_with_vars(y, *list(inputs), *[t.grad for t in optim.params])
+      if getenv("ASSIGN", 1): linear, var_vals = Tensor.linear_with_vars(y, *list(inputs), *optim.schedule_step())
+      else: linear, var_vals = Tensor.linear_with_vars(y, *list(inputs), *[t.grad for t in optim.params])
 
       for _ in range(JITCNT):
-        run_schedule(sched)
+        run_linear(linear, var_vals)
 
     CNT = getenv("CNT", 5)
     best_tm = None

test/external/process_replay/process_replay.py

@@ -8,8 +8,8 @@ ASSERT_DIFF = int((flag:="[pr]") in os.getenv("COMMIT_MESSAGE", flag) or flag in
 if not int(os.getenv("ASSERT_PROCESS_REPLAY", "1")): ASSERT_DIFF = 0
 
 try:
-  from tinygrad.renderer import Renderer, ProgramSpec
-  from tinygrad.engine.realize import get_program
+  from tinygrad.renderer import Renderer
+  from tinygrad.codegen import to_program
   from tinygrad.uop.ops import UOp, Ops
   from tinygrad.helpers import VERSION, Context, ContextVar, colored, db_connection, getenv, tqdm
 except ImportError as e:
@@ -41,24 +41,25 @@ class ProcessReplayWarning(Warning): pass
 
 # *** replay the function and convert return values to string
 
-def replay_get_program(p:ProgramSpec, ast:UOp, renderer:Renderer) -> tuple[str, str, tuple[Any, ...]]:
+def replay_to_program(p:UOp, ast:UOp, renderer:Renderer) -> tuple[str, str, tuple[Any, ...]]:
   if ast.op is Ops.PROGRAM: input_ast = ast
   else:
-    sink = ast.src[0] if ast.op is Ops.BEAM else ast
-    sink_arg = sink.arg
-    if ast.op is Ops.BEAM: sink_arg = replace(sink_arg, opts_to_apply=p.applied_opts)
-    input_ast = sink.replace(arg=replace(sink_arg, name=p.name))
-  p2 = get_program(input_ast, renderer=renderer)
-  def to_str(ret:ProgramSpec) -> str:
+    sink_arg = ast.arg
+    if sink_arg.beam: sink_arg = replace(sink_arg, opts_to_apply=p.src[0].arg.applied_opts)
+    input_ast = ast.replace(arg=replace(sink_arg, name=p.src[0].arg.name))
+  p2 = to_program(input_ast, renderer=renderer)
+  device = p.src[1].arg
+  def to_str(ret:UOp) -> str:
+    src = ret.src[3].arg
     # PYTHON renderer pickles UOps, first unpickle and decode here
-    if p.device.startswith("PYTHON"): return "\n".join([str(x) for x in pickle.loads(base64.b64decode(ret.src))])
-    return ret.src
+    if device.startswith("PYTHON"): return "\n".join([str(x) for x in pickle.loads(base64.b64decode(src))])
+    return src
   # properly color the name arg
   ast_repr = codecs.decode(str(input_ast), "unicode_escape")
   return to_str(p2), to_str(p), (ast_repr, renderer)
 
 replayers: dict[str, Callable[..., tuple[str, str, tuple[Any, ...]]]] = {}
-replayers["get_program"] = replay_get_program
+replayers["do_to_program"] = replay_to_program
 
 # *** run replayers on captured rows and print diffs
 

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/amd/README

@@ -20,17 +20,17 @@ test_llvm.py tests asm/disasm on the LLVM tests, confirming it behaves the same
 
 tinygrad's dtype tests should pass with and without LLVM. they run in about 12 seconds.
 
-`DEV=MOCKKFD+AMD pytest -n=12 test/backend/test_dtype_alu.py test/backend/test_dtype.py`
-`DEV=MOCKKFD+AMD:LLVM pytest -n=12 test/backend/test_dtype_alu.py test/backend/test_dtype.py`
+`DEV=MOCK+AMD pytest -n=12 test/backend/test_dtype_alu.py test/backend/test_dtype.py`
+`DEV=MOCK+AMD:LLVM pytest -n=12 test/backend/test_dtype_alu.py test/backend/test_dtype.py`
 
 The ops tests also pass, but they are very slow, so you should run them one at a time.
 
-`SKIP_SLOW_TEST=1 DEV=MOCKKFD+AMD pytest -n=12 test/backend/test_ops.py`
-`SKIP_SLOW_TEST=1 DEV=NOCKKFD+AMD:LLVM pytest -n=12 test/backend/test_ops.py`
+`SKIP_SLOW_TEST=1 DEV=MOCK+AMD pytest -n=12 test/backend/test_ops.py`
+`SKIP_SLOW_TEST=1 DEV=NOCK+AMD:LLVM pytest -n=12 test/backend/test_ops.py`
 
 When something is caught by main tinygrad tests, a local regression test should be added to `test/amd`.
 While working with tinygrad, you can dump the assembly with `DEBUG=7`. These tests all pass on real hardware
-If a test is failing with `DEV=MOCKKFD+AMD` it's because an instruction is emulated incorrectly.
+If a test is failing with `DEV=MOCK+AMD` it's because an instruction is emulated incorrectly.
 You can test with just `DEV=AMD` to test on real hardware, if it works on real hardware there's a bug in the emulator.
 IMPORTANT: if a test is failing in the emulator, it's an instruction bug. Use DEBUG=7, get the instructions, and debug.
 

test/mockgpu/amd/emu.py

@@ -161,7 +161,7 @@ def _init_sqtt_encoder():
     """Emit an SQTT packet for one executed instruction."""
     w = wave_id & 0x1F
     if wave_id not in started:
-      _emit_nibbles(nibbles, WAVESTART, delta=1, simd=0, cu_lo=0, wave=w, id7=wave_id)
+      _emit_nibbles(nibbles, WAVESTART, delta=1, simd=0, wgp=0, wave=w, id7=wave_id)
       started.add(wave_id)
     inst_type, inst_op, op_name = type(inst), inst.op.value if hasattr(inst, 'op') else 0, inst.op.name if hasattr(inst, 'op') else ""
     if issubclass(inst_type, _SOPP):
@@ -180,7 +180,7 @@ def _init_sqtt_encoder():
 
   def finish(wave_id: int):
     """Emit WAVEEND for a completed wave."""
-    if wave_id in started: _emit_nibbles(nibbles, WAVEEND, delta=1, simd=0, cu_lo=0, wave=wave_id & 0x1F)
+    if wave_id in started: _emit_nibbles(nibbles, WAVEEND, delta=1, simd=0, wgp=0, wave=wave_id & 0x1F)
 
   def finalize() -> bytes:
     """Pad and return the encoded SQTT blob."""

test/mockgpu/cuda/cuda.py

@@ -169,3 +169,7 @@ def cuGetErrorString(error: int, pStr) -> int:
   # Set the pointer to point to our error string buffer
   pStr._obj.value = ctypes.cast(buf, ctypes.POINTER(ctypes.c_char))
   return orig_cuda.CUDA_SUCCESS
+
+def cuDeviceGetCount(count) -> int:
+  count._obj.value = 1
+  return orig_cuda.CUDA_SUCCESS

test/mockgpu/mockgpu.py

@@ -11,8 +11,8 @@ libc = ctypes.CDLL(ctypes.util.find_library("c"))
 libc.mmap.argtypes = [ctypes.c_void_p, ctypes.c_size_t, ctypes.c_int, ctypes.c_int, ctypes.c_int, ctypes.c_long]
 libc.mmap.restype = ctypes.c_void_p
 
-drivers = [cls() for t in DEV.value if (cls:={"MOCKPCI+AMD": AMDriver, "MOCKKFD+AMD": AMDDriver, "MOCKUSB+AMD": AMUSBDriver,
-                                              "MOCKNVK+NV": NVDriver}.get(f"{t.interface}+{t.device}"))]
+drivers = [cls() for t in DEV.value if (cls:={"MOCKPCI+AMD": AMDriver, "MOCKKFD+AMD": AMDDriver, "MOCK+AMD": AMDDriver, "MOCKUSB+AMD": AMUSBDriver,
+                                              "MOCK+NV": NVDriver}.get(f"{t.interface}+{t.device}"))]
 tracked_fds = {}
 
 original_memoryview = builtins.memoryview
@@ -67,7 +67,7 @@ class MockFileIOInterface(FileIOInterface):
     if self.fd in tracked_fds:
       tracked_fds[self.fd].close(self.fd)
       tracked_fds.pop(self.fd)
-    else: os.close(self.fd)
+    elif self.fd is not None: os.close(self.fd)
 
   def ioctl(self, request, arg):
     if self.fd in tracked_fds:

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}"
@@ -128,7 +128,7 @@ class TestBitcastConstFolding(unittest.TestCase):
   def test_vec_bitcast(self):
     with Context(SPEC=0):
       r = full_rewrite_to_sink(UOp.const(dtypes.int32.vec(3), (-1, -2**31, 75)).bitcast(dtypes.uint32.vec(3)).sink()).src[0]
-    self.assertEqual(r.op, Ops.VECTORIZE)
+    self.assertEqual(r.op, Ops.STACK)
     self.assertEqual(r.dtype, dtypes.uint32.vec(3))
     self.assertEqual(tuple(x.arg for x in r.src), (2**32-1, 2**31, 75))
 

test/null/test_device.py

@@ -38,6 +38,13 @@ class TestDevice(unittest.TestCase):
     self.assertNotEqual(result.returncode, 0)
     self.assertIn(b"did you mean: 'USB'", result.stderr)
 
+  @unittest.skipIf(Device.DEFAULT != "AMD", "only run on AMD")
+  def test_dev_id_out_of_range(self):
+    result = subprocess.run(['python3', '-c', 'from tinygrad import Device; Device[Device.DEFAULT]'],
+                            env={**os.environ, "DEV":":99+AMD"}, capture_output=True)
+    self.assertNotEqual(result.returncode, 0)
+    self.assertIn(b"invalid visibility filter", result.stderr)
+
   def test_lowercase_canonicalizes(self):
     device = Device.DEFAULT
     with Context(DEV=device.lower()):
@@ -125,6 +132,7 @@ class TestDevVar(unittest.TestCase):
     for d, t in [("AMD", Target(device="AMD", renderer="")), ("AMD:LLVM", Target(device="AMD", renderer="LLVM")),
                  (":LLVM", Target(device="", renderer="LLVM")), ("AMD::gfx1100", Target(device="AMD", arch="gfx1100")),
                  ("AMD:LLVM:gfx1100", Target(device="AMD", renderer="LLVM", arch="gfx1100")), ("::gfx1100", Target(arch="gfx1100")),
+                 ("CPU:LLVM:arm64,native,AMX", Target(device="CPU", renderer="LLVM", arch="arm64,native,AMX")),
                  ("USB+", Target(interface="USB")), ("USB+AMD", Target(device="AMD", interface="USB")),
                  ("PCI:0+AMD", Target(device="AMD", interface="PCI", indices="0")), (":0+AMD", Target(device="AMD", indices="0")),
                  ("PCI:0,1+AMD", Target(device="AMD", interface="PCI", indices="0,1")),
@@ -145,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):