cleaner

* move cli

* update imports

* cleanup the readme

* edit

* work

* details

* python -m tinygrad.viz.cli

* do not execv in non tty

* option

* lint

* simpler

* gemm pmc

.github/workflows/autogen.yml

@@ -58,6 +58,7 @@ jobs:
         python3 -c "from tinygrad.runtime.autogen import avcodec"
         python3 -c "from tinygrad.runtime.autogen import llvm_qcom"
         python3 -c "from tinygrad.runtime.autogen import mlx5"
+        python3 -c "from tinygrad.runtime.autogen import ggml_common"
         REGEN=1 python3 -c "from tinygrad.runtime.autogen import libclang"
     - name: Check for differences
       run: |

.github/workflows/benchmark.yml

@@ -626,7 +626,7 @@ jobs:
     - 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
-      run: BENCHMARK_LOG=openpilot_0_11_0_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=3 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_policy.onnx
+      run: BENCHMARK_LOG=openpilot_0_11_0_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=4 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_policy.onnx
     - name: openpilot compile3 0.11.0 dmonitoring
       run: BENCHMARK_LOG=openpilot_0_11_0_dmonitoring PYTHONPATH="." ASSERT_MIN_STEP_TIME=11 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/dmonitoring_model.onnx
     - name: DEBUG=2 openpilot compile3 0.10.1 driving_vision
@@ -634,7 +634,7 @@ jobs:
     - name: openpilot compile3 0.10.1 driving_vision
       run: BENCHMARK_LOG=openpilot_0_10_1_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/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
     - name: openpilot compile3 0.10.1 driving_policy
-      run: BENCHMARK_LOG=openpilot_0_10_1_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=3 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_policy.onnx
+      run: BENCHMARK_LOG=openpilot_0_10_1_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=4 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_policy.onnx
     - name: openpilot compile3 0.10.1 dmonitoring
       run: BENCHMARK_LOG=openpilot_0_10_1_dmonitoring PYTHONPATH="." ASSERT_MIN_STEP_TIME=11 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/dmonitoring_model.onnx
     - name: benchmark MobileNetV2 on DSP

.github/workflows/test.yml

@@ -1,7 +1,7 @@
 name: Unit Tests
 env:
   # increment this when downloads substantially change to avoid the internet
-  CACHE_VERSION: '18'
+  CACHE_VERSION: '19'
   CAPTURE_PROCESS_REPLAY: 1
   GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
   PYTHONPATH: ${{ github.workspace }}
@@ -505,14 +505,14 @@ jobs:
         with:
           key: apps_llm
       - name: Test 1B LLM (llama)
-        run: echo "What's a male chicken called? Answer with only one word." | MAX_BUFFER_SIZE=0 python3 -m tinygrad.apps.llm --model llama3.2:1b | tee /dev/stderr | grep -i rooster
+        run: echo "What's a male chicken called? Answer with only one word." | MAX_BUFFER_SIZE=0 python3 -m tinygrad.llm --model llama3.2:1b | tee /dev/stderr | grep -i rooster
       - name: Test 1B LLM (llama q4)
-        run: echo "What's a male chicken called? Answer with only one word." | MAX_BUFFER_SIZE=0 python3 -m tinygrad.apps.llm --model llama3.2:1b-q4 | tee /dev/stderr | grep -i rooster
+        run: echo "What's a male chicken called? Answer with only one word." | MAX_BUFFER_SIZE=0 python3 -m tinygrad.llm --model llama3.2:1b-q4 | tee /dev/stderr | grep -i rooster
       - name: Test 1B LLM (qwen3.5)
-        run: echo "What's a male chicken called? Answer with only one word." | MAX_BUFFER_SIZE=0 python3 -m tinygrad.apps.llm --model qwen3.5:0.8b | tee /dev/stderr | grep -i rooster
+        run: echo "What's a male chicken called? Answer with only one word." | MAX_BUFFER_SIZE=0 python3 -m tinygrad.llm --model qwen3.5:0.8b | tee /dev/stderr | grep -i rooster
       - name: Test 1B LLM (qwen)
         # NOTE: qwen is dumb and only knows about female chickens
-        run: echo "What's a female chicken called? Answer with only one word." | MAX_BUFFER_SIZE=0 python3 -m tinygrad.apps.llm --model qwen3:0.6b | tee /dev/stderr | grep -i hen
+        run: echo "What's a female chicken called? Answer with only one word." | MAX_BUFFER_SIZE=0 python3 -m tinygrad.llm --model qwen3:0.6b | tee /dev/stderr | grep -i hen
 
 # ****** Models Tests ******
 
@@ -643,8 +643,7 @@ jobs:
     runs-on: ubuntu-24.04
     timeout-minutes: 20
     env:
-      DEV: AMD
-      MOCKGPU: 1
+      DEV: MOCKKFD+AMD
     steps:
       - name: Checkout Code
         uses: actions/checkout@v6
@@ -670,7 +669,7 @@ jobs:
       - name: Run AMD renderer tests
         run: python -m pytest -n=auto test/amd/ --durations 20
       - name: Run AMD renderer tests (AMD:LLVM)
-        run: DEV=AMD:LLVM python -m pytest -n=auto test/amd/ --durations 20
+        run: DEV=MOCKKFD+AMD:LLVM python -m pytest -n=auto test/amd/ --durations 20
       - name: Run SQTT profiling tests
         run: PROFILE=1 SQTT=1 python3 -m pytest -n=auto test/amd/test_sqtt_profiler.py
       - name: Run AMD emulated tests on NULL backend
@@ -679,20 +678,19 @@ jobs:
         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
-      - name: Run matmul on MOCKGPU
+      - name: Run matmul on MOCKKFD
         run: |
-          PYTHONPATH="." DEV=AMD MOCKGPU=1 N=256 python3 extra/gemm/amd_asm_matmul.py
-          PYTHONPATH="." DEV=AMD MOCKGPU=1 N=256 python3 extra/gemm/amd_copy_matmul.py
+          PYTHONPATH="." DEV=MOCKKFD+AMD N=256 python3 extra/gemm/amd_asm_matmul.py
+          PYTHONPATH="." DEV=MOCKKFD+AMD N=256 python3 extra/gemm/amd_copy_matmul.py
       - name: Run LLVM test
-        run: DEV=AMD:LLVM python test/device/test_amd_llvm.py
+        run: DEV=MOCKKFD+AMD:LLVM python test/device/test_amd_llvm.py
 
   testmockam:
     name: Linux (am)
     runs-on: ubuntu-24.04
     timeout-minutes: 15
     env:
-      DEV: PCI+AMD
-      MOCKGPU: 1
+      DEV: MOCKPCI+AMD
     steps:
       - name: Checkout Code
         uses: actions/checkout@v6
@@ -704,13 +702,13 @@ jobs:
           amd: 'true'
       - name: Run test_tiny on MOCKAM
         run: python test/test_tiny.py
-      - name: Run test_tiny on MOCKAM USB
-        run: GMMU=0 DEV=USB+AMD python test/test_tiny.py
-      - name: Run test_hcq on MOCKAM
+      - name: Run test_tiny on MOCKUSB
+        run: GMMU=0 DEV=MOCKUSB+AMD python test/test_tiny.py
+      - name: Run test_hcq on MOCKPCI
         run: python -m pytest test/device/test_hcq.py
-      - name: Run disk copy tests on MOCKAM
+      - name: Run disk copy tests on MOCKPCI
         run: python -m pytest test/unit/test_disk_tensor.py -k test_copy_from_disk
-      - name: Run test_tiny on MOCKAM Remote
+      - name: Run test_tiny on MOCKPCI Remote
         run: |
           python extra/remote/serve.py 6667 &
           sleep 2
@@ -728,8 +726,7 @@ jobs:
     runs-on: ubuntu-22.04
     timeout-minutes: 15
     env:
-      DEV: AMD:${{ matrix.backend == 'amdllvm' && 'LLVM' || '' }}:${{ matrix.arch }}
-      MOCKGPU: 1
+      DEV: MOCKKFD+AMD:${{ matrix.backend == 'amdllvm' && 'LLVM' || '' }}:${{ matrix.arch }}
       SKIP_SLOW_TEST: 1
     steps:
       - name: Checkout Code
@@ -764,7 +761,6 @@ jobs:
     runs-on: ubuntu-22.04
     timeout-minutes: 20
     env:
-      MOCKGPU: 1
       FORWARD_ONLY: 1
     steps:
       - name: Checkout Code
@@ -777,7 +773,7 @@ jobs:
           cuda: 'true'
           ocelot: 'true'
       - name: Set env
-        run: printf "${{ matrix.backend == 'ptx' && 'DEV=CUDA:PTX' || matrix.backend == 'nv' && 'DEV=NV\nSKIP_SLOW_TEST=1' }}" >> $GITHUB_ENV
+        run: printf "${{ matrix.backend == 'ptx' && 'DEV=MOCK+CUDA:PTX' || matrix.backend == 'nv' && 'DEV=MOCKNVK+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"
@@ -862,22 +858,19 @@ jobs:
       run: DEV=METAL TRANSCENDENTAL=2 python -m pytest -n=auto test/backend/test_ops.py::TestOps::test_sin test/backend/test_ops.py::TestOps::test_cos test/backend/test_ops.py::TestOps::test_tan test/backend/test_ops.py::TestOps::test_exp test/backend/test_ops.py::TestOps::test_log --durations=20
     - name: Run pytest (amd)
       env:
-        MOCKGPU: 1
-        DEV: AMD
+        DEV: MOCKKFD+AMD
         FORWARD_ONLY: 1
       run: |
         python3 -m pytest -n=auto test/device/test_hcq.py test/test_tiny.py --durations=20
     - name: Run pytest (amd with llvm backend)
       env:
-        MOCKGPU: 1
-        DEV: "AMD:LLVM"
+        DEV: "MOCKKFD+AMD:LLVM"
         FORWARD_ONLY: 1
       run: |
         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:
-        MOCKGPU: 1
-        DEV: "NV:PTX"
+        DEV: "MOCKNVK+NV:PTX"
         FORWARD_ONLY: 1
         # TODO: failing due to library loading error
         CAPTURE_PROCESS_REPLAY: 0
@@ -952,8 +945,8 @@ jobs:
       - name: Run process replay tests
         uses: ./.github/actions/process-replay
       - name: Run macOS-specific unit test
-        if: matrix.backend == 'cpu'
-        run: python3 -m pytest test/unit/test_disk_tensor.py::TestDiskTensor::test_copy_to_cpu_not_truncated
+        if: matrix.backend == 'llvm'
+        run: python3 -m pytest test/unit/test_disk_tensor.py::TestDiskTensor::test_copy_to_cpu_not_truncated test/unit/test_cpu.py
 
 # ****** Windows Tests ******
 

docs/abstractions4.py

@@ -1,9 +1,9 @@
 # tinygrad allows you to write kernels at many different abstractions levels.
 # This is for RDNA3, but if you don't have one you can run with the emulator
-# PYTHONPATH="." MOCKGPU=1 DEV=AMD
+# PYTHONPATH="." DEV=MOCKPCI+AMD
 
 from tinygrad import Tensor, Context, GlobalCounters, UOp, Device
-from tinygrad.helpers import DEBUG, getenv
+from tinygrad.helpers import DEV, DEBUG, getenv
 from tinygrad.uop.ops import AxisType, KernelInfo, Ops
 from tinygrad.dtype import AddrSpace, dtypes
 from tinygrad.runtime.autogen.amd.rdna3.ins import *
@@ -16,7 +16,7 @@ def eval_harness(name, tensor, fxn, check=None):
   print(f"computed in {GlobalCounters.time_sum_s*1000:.2f} ms, {(a.nbytes()/1e9)/GlobalCounters.time_sum_s:.2f} GB/s")
   return out
 
-SZ = 256*1024 if getenv("MOCKGPU") else 1024*1024*1024
+SZ = 256*1024 if DEV.interface.startswith("MOCK") else 1024*1024*1024
 
 def example_2_hip(a:Tensor, correct):
   GLOBALS = 1024

docs/nn.md

@@ -37,4 +37,4 @@
     options:
         show_signature: false
         separate_signature: false
-::: tinygrad.nn.state.gguf_load
+::: tinygrad.llm.gguf.gguf_load

docs/tinygpu.md

@@ -55,7 +55,7 @@ export PATH="$HOME/.local/bin:$PATH"
 ### 5. Use it!
 
 ```bash
-DEV={AMD|NV} python3 tinygrad/apps/llm.py
+DEV={AMD|NV} python3 -m tinygrad.llm
 ```
 
 **Note:** Use `JITBEAM=2` to search for faster kernels (one-time search cost, results cached).

examples/compile_tensorflow.py

@@ -35,12 +35,11 @@ def compile_onnx_model(onnx_model):
   tinyonnx = TinyOnnx(onnx_model)
   the_input = Tensor.randn(1,32)
 
-  run, special_names = jit_model(tinyonnx, the_input)
+  linear, output_bufs = jit_model(tinyonnx, the_input)
+  the_output = [tinyonnx.forward(the_input)]
 
-  functions, statements, bufs, bufs_to_save = compile_net(run, special_names)
+  functions, statements, bufs, bufs_to_save = compile_net(linear, output_bufs)
   prg = export_model_clang(functions, statements, bufs, {}, ["input0"], ["output0"])
-
-  the_output = run(the_input)
   cprog = ["#include <string.h>", "#include <stdio.h>", "#include <stdlib.h>"]
   cprog.append(prg)
 

examples/gpt2.py

@@ -5,8 +5,9 @@ with contextlib.suppress(ImportError): import tiktoken
 from tinygrad import Tensor, TinyJit, Device, GlobalCounters, Variable, dtypes
 from tinygrad.uop.ops import UOp
 from tinygrad.helpers import Timing, DEBUG, JIT, getenv, fetch, colored, trange
+from tinygrad.llm.gguf import gguf_load
 from tinygrad.nn import Embedding, Linear, LayerNorm
-from tinygrad.nn.state import gguf_load, torch_load, load_state_dict, get_state_dict
+from tinygrad.nn.state import torch_load, load_state_dict, get_state_dict
 from extra.bench_log import BenchEvent, WallTimeEvent
 
 MAX_CONTEXT = getenv("MAX_CONTEXT", 128)

examples/llama3.py

@@ -2,7 +2,8 @@ from pathlib import Path
 from typing import List
 import json, argparse, random, time, os
 from extra.models.llama import Transformer, convert_from_huggingface, convert_from_gguf, fix_bf16
-from tinygrad.nn.state import safe_load, torch_load, load_state_dict, get_parameters, gguf_load
+from tinygrad.llm.gguf import gguf_load
+from tinygrad.nn.state import safe_load, torch_load, load_state_dict, get_parameters
 from tinygrad import Tensor, dtypes, nn, Context, Device, GlobalCounters
 from tinygrad.helpers import Profiling, Timing, DEBUG, colored, fetch, tqdm
 from extra.bench_log import BenchEvent, WallTimeEvent

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
+  from examples.mlperf.models.flat_llama import FlatTransformer, apply_grad, FP8, FP8_DTYPE
   from examples.llama3 import MODEL_PARAMS
   from examples.mlperf.lr_schedulers import CosineAnnealingLRWithWarmup
   from examples.mlperf.optim import GradAccClipAdamW
@@ -1416,9 +1416,9 @@ def train_llama3():
   optim = GradAccClipAdamW(params, lr=0.0, b1=opt_adamw_beta_1, b2=opt_adamw_beta_2,
                            eps=opt_adamw_epsilon, weight_decay=opt_adamw_weight_decay, grad_acc=grad_acc, device=optim_device)
 
-  # init grads
   for p in optim.params:
-    p.grad = Tensor.zeros(p.shape, dtype=p.dtype, device=p.device).contiguous()
+    grad_dtype = dtypes.bfloat16 if p.dtype == FP8_DTYPE else p.dtype
+    p.grad = Tensor.zeros(p.shape, dtype=grad_dtype, device=p.device).contiguous()
   grads = [p.grad for p in optim.params]
 
   scheduler = CosineAnnealingLRWithWarmup(optim, opt_base_learning_rate, opt_end_learning_rate, opt_learning_rate_warmup_steps, opt_learning_rate_decay_steps)
@@ -1433,6 +1433,17 @@ def train_llama3():
     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 []
+
+  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())
 
   @TinyJit
   def minibatch(tokens:Tensor):
@@ -1457,7 +1468,7 @@ def train_llama3():
 
     lr_cpu = optim.lr.float().to("CPU")
     grad_norm_cpu = grad_norm.float().to("CPU")
-    Tensor.realize(lr_cpu, grad_norm_cpu, *grads)
+    Tensor.realize(lr_cpu, grad_norm_cpu, *grads, *fp8_inv_scales)
 
     return lr_cpu, grad_norm_cpu
 

examples/mlperf/models/flat_llama.py

@@ -42,19 +42,19 @@ def quantize_fp8(x:Tensor, amax_state:Tensor|None=None):
   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, amax_w:Tensor|None=None) -> tuple[Tensor,...]:
+def matmul(x:Tensor, w:Tensor, fp8=FP8, 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"):
       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,)
-  x_fp8, x_scale, x_new_amax = quantize_fp8(x, amax_state=amax_x)
-  w_fp8, w_scale, w_new_amax = quantize_fp8(w, amax_state=amax_w)
-  combined_scale = x_scale * w_scale
+  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"):
     from extra.gemm.cdna_asm_gemm import can_use_asm_gemm, asm_gemm
-    if can_use_asm_gemm(x_fp8, w_fp8.T): return asm_gemm(x_fp8, w_fp8.T, combined_scale=combined_scale), x_new_amax, w_new_amax, x_fp8, w_fp8
-  return x_fp8.dot(w_fp8.T, dtype=dtypes.float) * combined_scale, x_new_amax, w_new_amax, x_fp8, w_fp8
+    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
 
 def _rmsnorm_fwd(x_in:Tensor, eps:float) -> tuple[Tensor, Tensor]:
   x = x_in.float()
@@ -85,17 +85,18 @@ class FlatTransformer:
     self.n_kv_heads = n_kv_heads if n_kv_heads is not None else n_heads # n_kv_heads != n_heads implies MQA [arxiv/2307.09288, A.2.1]
     self.head_dim = dim // n_heads
     self.n_rep = self.n_heads // self.n_kv_heads
+    self.hidden_dim = hidden_dim
 
     scaled_std = 0.02 / math.sqrt(2 * n_layers)
 
     # Attention
+    self._init_inv_scales = []  # populated by lin_per_layer when FP8
     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)
 
     # FeedForward
-    self.w1 = self.lin_per_layer(dim, hidden_dim)
+    self.w13 = self.lin_per_layer(dim, hidden_dim * 2)
     self.w2 = self.lin_per_layer(hidden_dim, dim, std=scaled_std)
-    self.w3 = self.lin_per_layer(dim, hidden_dim)
 
     self.norm_eps = norm_eps
     self.attention_norm = Tensor.ones(n_layers, dim).contiguous()
@@ -110,18 +111,27 @@ class FlatTransformer:
 
     if FP8:
       def _amax(): return Tensor.full((), FP8_MAX).contiguous().requires_grad_(False)
-      names = ["xqkv", "wqkv", "xo", "wo", "x1", "w1", "x2", "w2", "x3", "w3"]
-      # _fp8_amax[name][layer_idx] = scalar amax tensor
+      names = ["xqkv", "xo", "x13", "x2"]
       self._fp8_amax = {name: [_amax() for _ in range(n_layers)] for name in names}
-      self._fp8_amax["xout"] = [_amax()]
-      self._fp8_amax["wout"] = [_amax()]
+      # 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"): return Tensor.zeros(self.n_layers, out_features, in_features)
-    return Tensor.normal(self.n_layers, out_features, in_features, mean=0.0, std=std)
+    if getenv("ZEROS"): 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)
+    self._init_inv_scales.append((amax + 1e-8) / FP8_MAX)  # save for inv_scale init
+    return (w * scale.reshape(-1, 1, 1)).clamp(-FP8_MAX, FP8_MAX).cast(FP8_DTYPE)
 
   def attention(self, x:Tensor, freqs_cis:Tensor, attention_norm:Tensor, wqkv:Tensor, wo:Tensor,
-                amax_xqkv=None, amax_wqkv=None, amax_xo=None, amax_wo=None):
+                amax_xqkv=None, amax_xo=None, s_qkv=None, s_o=None):
     bsz, seqlen, _ = x.shape
     new_amaxs, saves = [], []
 
@@ -129,9 +139,9 @@ class FlatTransformer:
     saves.extend([x, rrms])
     x = x * attention_norm
 
-    xqkv, *ret = matmul(x, wqkv, amax_x=amax_xqkv, amax_w=amax_wqkv)
-    new_amaxs.extend(ret[:2])
-    saves.extend(ret[2:] + [xqkv])
+    xqkv, *ret = matmul(x, wqkv, amax_x=amax_xqkv, w_inv_scale=s_qkv)
+    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)
     xq = xqkv[:, :, :, :self.n_rep].reshape(bsz, seqlen, self.n_heads, self.head_dim)
     xk = xqkv[:, :, :, self.n_rep].reshape(bsz, seqlen, self.n_kv_heads, self.head_dim)
@@ -148,43 +158,51 @@ class FlatTransformer:
       attn = xq.scaled_dot_product_attention(xk, xv, is_causal=True, enable_gqa=True)
     attn = attn.transpose(1, 2).reshape(bsz, seqlen, -1)
 
-    out, *ret = matmul(attn, wo, amax_x=amax_xo, amax_w=amax_wo)
-    new_amaxs.extend(ret[:2])
-    saves.extend(ret[2:] + [out])
+    out, *ret = matmul(attn, wo, amax_x=amax_xo, w_inv_scale=s_o)
+    new_amaxs.extend(ret[:1])
+    saves.extend(ret[1:] + [out])
     return (out, *new_amaxs, *saves)
 
-  def feed_forward(self, x:Tensor, ffn_norm:Tensor, w1:Tensor, w2:Tensor, w3:Tensor,
-                   amax_x1=None, amax_w1=None, amax_x2=None, amax_w2=None, amax_x3=None, amax_w3=None):
+  def feed_forward(self, x:Tensor, ffn_norm:Tensor, w13:Tensor, w2:Tensor,
+                   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_w1, *ret = matmul(x, w1, amax_x=amax_x1, amax_w=amax_w1)
-    new_amaxs.extend(ret[:2])
-    saves.extend(ret[2:] + [x_w1])
-    x_w3, *ret = matmul(x.contiguous_backward(), w3, amax_x=amax_x3, amax_w=amax_w3)
-    new_amaxs.extend(ret[:2])
-    saves.extend(ret[2:] + [x_w3])
-    out, *ret = matmul(x_w1.silu() * x_w3, w2, amax_x=amax_x2, amax_w=amax_w2)
-    new_amaxs.extend(ret[:2])
-    saves.extend(ret[2:] + [out])
+    x_w13, *ret = matmul(x, w13, amax_x=amax_x13, w_inv_scale=s_13)
+    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)
+    new_amaxs.extend(ret[:1])
+    saves.extend(ret[1:] + [out])
     return (out, *new_amaxs, *saves)
 
   @function(precompile=True, precompile_backward=True)
   def run_layer(self, x:Tensor, freqs_cis:Tensor,
                 attention_norm:Tensor, wqkv:Tensor, wo:Tensor,
-                ffn_norm:Tensor, w1:Tensor, w2:Tensor, w3:Tensor,
-                amax_xqkv=None, amax_wqkv=None, amax_xo=None, amax_wo=None,
-                amax_x1=None, amax_w1=None, amax_x2=None, amax_w2=None, amax_x3=None, amax_w3=None):
+                ffn_norm:Tensor, w13:Tensor, w2:Tensor,
+                amax_xqkv=None, amax_xo=None,
+                amax_x13=None, amax_x2=None,
+                s_qkv=None, s_o=None, s_13=None, s_2=None):
     attn, *attn_ret = self.attention(x, freqs_cis, attention_norm, wqkv, wo,
-                                     amax_xqkv=amax_xqkv, amax_wqkv=amax_wqkv, amax_xo=amax_xo, amax_wo=amax_wo)
-    attn_amaxs, attn_saves = attn_ret[:4], attn_ret[4:]
+                                     amax_xqkv=amax_xqkv, amax_xo=amax_xo,
+                                     s_qkv=s_qkv, s_o=s_o)
+    attn_amaxs, attn_saves = attn_ret[:2], attn_ret[2:]
     h = x + attn
-    ffn, *ffn_ret = self.feed_forward(h, ffn_norm, w1, w2, w3,
-                                      amax_x1=amax_x1, amax_w1=amax_w1, amax_x2=amax_x2, amax_w2=amax_w2, amax_x3=amax_x3, amax_w3=amax_w3)
-    ffn_amaxs, ffn_saves = ffn_ret[:6], ffn_ret[6:]
+    ffn, *ffn_ret = self.feed_forward(h, ffn_norm, w13, w2,
+                                      amax_x13=amax_x13, amax_x2=amax_x2,
+                                      s_13=s_13, s_2=s_2)
+    ffn_amaxs, ffn_saves = ffn_ret[:2], ffn_ret[2:]
     h = h + ffn
     return (h, *attn_amaxs, *ffn_amaxs, *attn_saves, *ffn_saves)
 
@@ -196,9 +214,8 @@ class FlatTransformer:
       # 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.w1.shard_(device, axis=1).realize()             # (n_layers, hidden, dim) shard out
+      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.w3.shard_(device, axis=1).realize()             # (n_layers, hidden, dim) shard out
       self.attention_norm.shard_(device, axis=None).realize()
       self.ffn_norm.shard_(device, axis=None).realize()
       self.norm.weight.shard_(device, axis=None).realize()
@@ -209,24 +226,26 @@ class FlatTransformer:
         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
     for i in range(self.n_layers):
-      amax_layer = {"amax_xqkv": a["xqkv"][i], "amax_wqkv": a["wqkv"][i],
-                    "amax_xo": a["xo"][i], "amax_wo": a["wo"][i],
-                    "amax_x1": a["x1"][i], "amax_w1": a["w1"][i],
-                    "amax_x2": a["x2"][i], "amax_w2": a["w2"][i],
-                    "amax_x3": a["x3"][i], "amax_w3": a["w3"][i]} if a else {}
+      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.w1[i], self.w2[i], self.w3[i],
-                               **amax_layer)
+                               self.ffn_norm[i], self.w13[i], self.w2[i],
+                               **amax_layer, **scale_layer)
       if a:
-        amaxs = ret[:10]
-        amax_names = ["xqkv", "wqkv", "xo", "wo", "x1", "w1", "x3", "w3", "x2", "w2"]
+        amaxs = ret[:5]
+        amax_names = ["xqkv", "xo", "x13", "x2"]
         for name, new_val in zip(amax_names, amaxs):
           a[name][i].assign(new_val)
 

examples/mlperf/optim.py

@@ -34,7 +34,9 @@ class GradAccClipAdamW(Optimizer):
     else:
       updates, extra = self._step([], grads)
     for i, tt in enumerate(self.params): tt.assign(self._apply_update(tt, updates[i], self.master_params[i] if self.master_params else None))
-    to_realize = extra+self.params+self.buffers+(self.master_params or [])
+    # collect inv_scale tensors attached to fp8 params (set by _apply_update)
+    fp8_inv_scales = [tt._inv_scale for tt in self.params if hasattr(tt, '_inv_scale')]
+    to_realize = extra+self.params+self.buffers+(self.master_params or [])+fp8_inv_scales
 
     Tensor.realize(*to_realize)
     return extra[-1]
@@ -77,4 +79,12 @@ class GradAccClipAdamW(Optimizer):
     new_w = w.detach() - up
     if master is not None: master.assign(new_w)
     if STOCHASTIC_ROUND and t.dtype == dtypes.bfloat16: return stochastic_round_bf16(new_w)
+    if t.dtype in dtypes.fp8s:
+      from examples.mlperf.models.flat_llama import FP8_MAX
+      amax = new_w.float().abs().flatten(1).max(1).detach()  # per-layer amax for (n_layers, out, in)
+      scale = FP8_MAX / (amax + 1e-8)
+      fp8_w = (new_w * scale.reshape(-1, *([1]*(new_w.ndim-1)))).clamp(-FP8_MAX, FP8_MAX).cast(t.dtype)
+      if hasattr(t, '_inv_scale'):
+        t._inv_scale.assign(((amax + 1e-8) / FP8_MAX).cast(t._inv_scale.dtype))
+      return fp8_w
     return new_w.cast(t.dtype)

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"
-extra/viz/cli.py --profile -s "$SRC"
+python -m tinygrad.viz.cli --profile -s "$SRC" --top 20

examples/openpilot/compile3.py

@@ -4,7 +4,8 @@ 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.engine.realize import CompiledRunner
+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"
@@ -35,7 +36,11 @@ def compile(onnx_file):
       ret = run_onnx_jit(**inputs).numpy()
     # copy i == 1 so use of JITBEAM is okay
     if i == 1: test_val = np.copy(ret)
-  print(f"captured {len(run_onnx_jit.captured.jit_cache)} kernels")
+  # iterate kernel CALLs in the captured LINEAR UOp; toposort descends into batched graph CUSTOM_FUNCTIONs
+  kernel_asts = {Ops.SINK, Ops.PROGRAM, Ops.BEAM}
+  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")
   np.testing.assert_equal(test_val, ret, "JIT run failed")
   print("jit run validated")
 
@@ -43,13 +48,14 @@ def compile(onnx_file):
   kernel_count = 0
   read_image_count = 0
   gated_read_image_count = 0
-  for ei in run_onnx_jit.captured.jit_cache:
-    if isinstance(ei.prg, CompiledRunner):
-      kernel_count += 1
-      read_image_count += ei.prg.p.src.count("read_image")
-      gated_read_image_count += ei.prg.p.src.count("?read_image")
-      for v in [m.group(1) for m in re.finditer(r'(val\d+)\s*=\s*read_imagef\(', ei.prg.p.src)]:
-        if len(re.findall(fr'[\?\:]{v}\.[xyzw]', ei.prg.p.src)) > 0: gated_read_image_count += 1
+  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
+    kernel_count += 1
+    read_image_count += src.count("read_image")
+    gated_read_image_count += src.count("?read_image")
+    for v in [m.group(1) for m in re.finditer(r'(val\d+)\s*=\s*read_imagef\(', src)]:
+      if len(re.findall(fr'[\?\:]{v}\.[xyzw]', src)) > 0: gated_read_image_count += 1
   print(f"{kernel_count=},  {read_image_count=}, {gated_read_image_count=}")
   if (allowed_kernel_count:=getenv("ALLOWED_KERNEL_COUNT", -1)) != -1:
     assert kernel_count == allowed_kernel_count, f"different kernels! {kernel_count=}, {allowed_kernel_count=}"

examples/webgpu/stable_diffusion/compile.py

@@ -111,19 +111,19 @@ if __name__ == "__main__":
     return code
 
   def compile_step(model, step: Step):
-    run, special_names = jit_model(step, *step.input)
-    functions, statements, bufs, _ = compile_net(run, special_names)
+    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()}
     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 = [name for _,name in special_names.items() if "input" in name]
-    output_names = [name for _,name in special_names.items() if "output" in name]
+    input_names = [f"input{i}" for i in range(len(step.input))]
+    output_names = [f"output{i}" for i in range(len(output_bufs))]
     input_buf_types = [dtype_to_js_type(bufs[inp_name][1]) for inp_name in input_names]
     output_buf_types = [dtype_to_js_type(bufs[out_name][1]) for out_name in output_names]
     kernel_calls = '\n        '.join([f"addComputePass(device, commandEncoder, piplines[{i}], [{', '.join(args)}], {global_size});" for i, (_name, args, global_size, _local_size) in enumerate(statements) ])
     exported_bufs =  '\n    '.join([f"const {name} = " + (f"createEmptyBuf(device, {size});" if _key not in weights else f"createWeightBuf(device, {size}, getTensorBuffer(safetensor, metadata['{weights[_key]}'], '{weights[_key]}'))") + ";"  for name,(size,dtype,_key) in bufs.items()])
-    gpu_write_bufs =  '\n    '.join([f"const gpuWriteBuffer{i} = device.createBuffer({{size:input{i}.size, usage: GPUBufferUsage.COPY_SRC | GPUBufferUsage.MAP_WRITE }});" for i,(_,value) in enumerate(special_names.items()) if "output" not in value])
+    gpu_write_bufs =  '\n    '.join([f"const gpuWriteBuffer{i} = device.createBuffer({{size:input{i}.size, usage: GPUBufferUsage.COPY_SRC | GPUBufferUsage.MAP_WRITE }});" for i in range(len(input_names))])
     input_writer = '\n    '.join([f"await gpuWriteBuffer{i}.mapAsync(GPUMapMode.WRITE);\n    new {input_buf_types[i]}(gpuWriteBuffer{i}.getMappedRange()).set(" + f'data{i});' + f"\n    gpuWriteBuffer{i}.unmap();\ncommandEncoder.copyBufferToBuffer(gpuWriteBuffer{i}, 0, input{i}, 0, gpuWriteBuffer{i}.size);"  for i,_ in enumerate(input_names)])
     return f"""\n    var {step.name} = function() {{
 
@@ -141,7 +141,7 @@ if __name__ == "__main__":
         const kernels = [{kernel_names}];
         const piplines = await Promise.all(kernels.map(name => device.createComputePipelineAsync({{layout: "auto", compute: {{ module: device.createShaderModule({{ code: name }}), entryPoint: "main" }}}})));
 
-        return async ({",".join([f'data{i}' for i,(k,v) in enumerate(special_names.items()) if v != "output0"])}) => {{
+        return async ({",".join([f'data{i}' for i in range(len(input_names))])}) => {{
             const commandEncoder = device.createCommandEncoder();
 
             {input_writer}

extra/amax/cast_amax.py

@@ -0,0 +1,85 @@
+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
+
+@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)
+  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=(*sink.src, sink)),
+                               UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=lib)))
+
+@functools.cache
+def _custom_fused_cast_amax_w13(fp8_out:UOp, amax_buf:UOp, xw13:UOp, amax_state:UOp, dname:str) -> UOp:
+  hidden = xw13.shape[2] // 2
+  n_elems = xw13.shape[0] * xw13.shape[1] * hidden
+  threads, workgroups = UOp.special(THREADS_PER_WG, "lidx0"), UOp.special(NUM_WG, "gidx0")
+  # 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)
+  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
+  _, _, 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)
+    dname = device[0].split(":")[0]
+  else:
+    grad_xw13 = Tensor.invalid(*xw13.shape, dtype=dtypes.bfloat16, device=device)
+    dname = device.split(":")[0] if isinstance(device, str) else device
+  grad_x2_t = Tensor(gradient, device=device).cast(dtypes.bfloat16)
+  fxn = functools.partial(_custom_fused_bwd_w13, dname=dname)
+  grad_xw13, *_ = Tensor.custom_kernel(grad_xw13, Tensor(xw13, device=device), grad_x2_t, Tensor(amax_state, device=device), fxn=fxn)
+  return (None, None, grad_xw13.uop, None)
+
+def fused_quantize_fp8_w13(xw13:Tensor, amax_state:Tensor, fp8_dtype) -> tuple[Tensor, Tensor, Tensor]:
+  # silu(xw1)*xw3 -> fp8 + amax over fused xw13 layout. Returns (fp8, inv_scale, new_amax).
+  assert xw13.dtype == dtypes.bfloat16, f"expected bf16, got {xw13.dtype}"
+  MBS, SEQ, H2 = xw13.shape
+  assert H2 % 2 == 0, f"w13 last-axis must be even, got {H2}"
+  HIDDEN = H2 // 2
+  if isinstance(xw13.device, tuple):
+    axis, ndev = xw13.uop.axis, len(xw13.device)
+    assert axis in (0, 1), f"unsupported sharding axis={axis}"
+    fp8_out = Tensor(Tensor.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)
+    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)
+    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

extra/amax/cast_amax_bwd_w13.cpp

@@ -0,0 +1,68 @@
+#include <hip/hip_runtime.h>
+#include <hip/hip_bf16.h>
+
+#ifndef N_ELEMS
+#define N_ELEMS 234881024
+#endif
+#ifndef HIDDEN
+#define HIDDEN 14336
+#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_silu_mul_bwd_w13(
+    __hip_bfloat16*       __restrict__ grad_xw13_out,    // bf16, 2*N_ELEMS (interleaved layout)
+    const __hip_bfloat16* __restrict__ xw13,             // bf16, 2*N_ELEMS (interleaved)
+    const __hip_bfloat16* __restrict__ grad_x2,          // bf16, N_ELEMS
+    const __hip_bfloat16* __restrict__ amax_state)       // bf16 scalar
+{
+  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);
+
+  for (int base = gid * VEC; base < N_ELEMS; base += stride_elems) {
+    const int outer = base / HIDDEN;
+    const int inner = base % HIDDEN;
+    const int xw1_off = outer * 2 * HIDDEN + inner;
+    const int xw3_off = xw1_off + HIDDEN;
+
+    float4 x1_raw = *reinterpret_cast<const float4*>(&xw13[xw1_off]);
+    float4 x3_raw = *reinterpret_cast<const float4*>(&xw13[xw3_off]);
+    float4 g_raw  = *reinterpret_cast<const float4*>(&grad_x2[base]);
+
+    const __hip_bfloat16 *x1 = reinterpret_cast<const __hip_bfloat16*>(&x1_raw);
+    const __hip_bfloat16 *x3 = reinterpret_cast<const __hip_bfloat16*>(&x3_raw);
+    const __hip_bfloat16 *gv = reinterpret_cast<const __hip_bfloat16*>(&g_raw);
+
+    __hip_bfloat16 out1[VEC], out3[VEC];
+    #pragma unroll
+    for (int i = 0; i < VEC; i++) {
+      const float f1 = static_cast<float>(x1[i]);
+      const float f3 = static_cast<float>(x3[i]);
+      const float fg = static_cast<float>(gv[i]);
+      const float sig = 1.0f / (1.0f + __expf(-f1));
+      const float silu = f1 * sig;
+      const float silu_prime = sig + silu * (1.0f - sig);
+      const float gs = fg * scale;
+      out1[i] = static_cast<__hip_bfloat16>(gs * silu_prime * f3);
+      out3[i] = static_cast<__hip_bfloat16>(gs * silu);
+    }
+
+    *reinterpret_cast<float4*>(&grad_xw13_out[xw1_off]) = *reinterpret_cast<float4*>(out1);
+    *reinterpret_cast<float4*>(&grad_xw13_out[xw3_off]) = *reinterpret_cast<float4*>(out3);
+  }
+}

extra/amax/cast_amax_fwd_w13.cpp

@@ -0,0 +1,79 @@
+#include <hip/hip_runtime.h>
+#include <hip/hip_bf16.h>
+#include <hip/hip_fp8.h>
+
+#ifndef N_ELEMS
+#define N_ELEMS 234881024
+#endif
+#ifndef HIDDEN
+#define HIDDEN 14336
+#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 (so VEC loads don't straddle block boundary)");
+
+extern "C" __global__ __launch_bounds__(THREADS_PER_WG) void
+fused_silu_mul_cast_amax_w13(
+    __hip_fp8_storage_t*  __restrict__ fp8_out,         // fp8, N_ELEMS
+    __hip_bfloat16*       __restrict__ amax_buf,        // bf16, NUM_WG (per-WG amaxes)
+    const __hip_bfloat16* __restrict__ xw13,            // bf16, 2*N_ELEMS
+    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;
+
+  // grid-stride over 8-element groups
+  for (int base = gid * VEC; base < N_ELEMS; base += stride_elems) {
+    // interleaved xw13 layout: xw1 and xw3 are not contiguous halves
+    const int outer = base / HIDDEN;
+    const int inner = base % HIDDEN;
+    const int xw1_off = outer * 2 * HIDDEN + inner;
+    const int xw3_off = xw1_off + HIDDEN;
+
+    float4 x1_raw = *reinterpret_cast<const float4*>(&xw13[xw1_off]);
+    float4 x3_raw = *reinterpret_cast<const float4*>(&xw13[xw3_off]);
+
+    const __hip_bfloat16 *x1 = reinterpret_cast<const __hip_bfloat16*>(&x1_raw);
+    const __hip_bfloat16 *x3 = reinterpret_cast<const __hip_bfloat16*>(&x3_raw);
+
+    __hip_fp8_storage_t out[VEC];
+    #pragma unroll
+    for (int i = 0; i < VEC; i++) {
+      const float f1 = static_cast<float>(x1[i]);
+      const float f3 = static_cast<float>(x3[i]);
+      const float silu = f1 / (1.0f + __expf(-f1));
+      const float x2 = silu * f3;
+      local_max = fmaxf(local_max, fabsf(x2));
+      const float x_scaled = fmaxf(-FP8_MAX, fminf(FP8_MAX, x2 * scale));
+      out[i] = __hip_cvt_float_to_fp8(x_scaled, __HIP_SATFINITE, __HIP_E4M3);
+    }
+
+    *reinterpret_cast<uint64_t*>(&fp8_out[base]) = *reinterpret_cast<uint64_t*>(out);
+  }
+
+  // LDS tree reduction: per-workgroup 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/export_model.py

@@ -1,47 +1,49 @@
 from typing import Tuple, Dict, List, Optional
 from tinygrad.dtype import DType, dtypes
-from tinygrad.renderer import ProgramSpec
 from tinygrad.tensor import Tensor
-from tinygrad.device import Device
+from tinygrad.device import Device, Buffer
 from tinygrad.engine.jit import TinyJit
 from tinygrad.nn.state import get_state_dict
-from tinygrad.helpers import Context, to_mv
-from tinygrad.uop.ops import Ops
+from tinygrad.helpers import Context, to_mv, prod
+from tinygrad.uop.ops import Ops, UOp
+from tinygrad.codegen import get_program
 import json
 from collections import OrderedDict
 
 EXPORT_SUPPORTED_DEVICE = ["WEBGPU", "CPU", "CUDA", "CL"]
 
-def compile_net(run:TinyJit, special_names:Dict[int,str]) -> Tuple[Dict[str,str],List[Tuple[str,List[str],List[int]]],Dict[str,Tuple[int,DType,int]],Dict[str,Tensor]]:
-  # memory-planned subbuffers can have multiple Buffer objects for the same memory region
-  canon, _seen = {}, {}
-  for ji in run.jit_cache:
-    for b in ji.bufs:
-      if b is not None: canon[id(b)] = _seen.setdefault((id(b.base._buf), b.offset, b.size, b.dtype), b)
-  special_names = {id(canon[k]): v for k, v in special_names.items() if k in canon}
+_KERNEL_ASTS = {Ops.SINK, Ops.PROGRAM, Ops.BEAM}
+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)
 
-  functions, bufs, bufs_to_save, statements, bufnum = {}, {}, {}, [], 0
-  for ji in run.jit_cache:
-    fxn: ProgramSpec = ji.prg.p
-    functions[fxn.function_name] = fxn.src   # NOTE: this assumes all with the same name are the same
-    cargs = []
-    for i,arg in enumerate(ji.bufs):
-      arg = canon[id(arg)]
-      key = id(arg)
-      if key not in bufs:
-        if key in special_names:
-          bufs[key] = (special_names[key], arg.size*arg.dtype.itemsize, arg.dtype, key)
-        else:
-          bufs[key] = (f"buf_{bufnum}", arg.size*arg.dtype.itemsize, arg.dtype, key)
-          bufnum += 1
-          if i > 0: bufs_to_save[bufs[key][0]] = arg   # if first usage of a buffer is not an output, and it's not a special name
-      cargs.append(bufs[key][0])
-    cargs += [var for var in fxn.vars if getattr(var, "op", None) is Ops.DEFINE_VAR] # symbolic vars; is it necessary or sufficient to check for DEFINE_VAR?
-    statements.append((fxn.function_name, cargs, fxn.global_size, fxn.local_size))
+def compile_net(linear:UOp, output_bufs:List[Buffer]) -> Tuple[Dict[str,str], List, Dict[str,Tuple[int,DType,int]], Dict[str,Buffer]]:
+  output_name = {id(b): f"output{i}" for i, b in enumerate(output_bufs)}
+  functions, bufs, bufs_to_save, statements, n = {}, {}, {}, [], 0
 
-  return functions, statements, {name:(size, dtype, key) for (name,size,dtype,key) in bufs.values()}, bufs_to_save
+  def name_of(bu:UOp, is_out:bool) -> str:
+    nonlocal n
+    if bu.op is Ops.PARAM: key, name, size = ("in", bu.arg), f"input{bu.arg}", prod(bu.shape)*bu.dtype.itemsize
+    else:
+      b = bu.buffer
+      key, size = (id(b.base), b.offset, b.size, b.dtype), b.size*b.dtype.itemsize
+      if key in bufs: return bufs[key][0]
+      if (name:=output_name.get(id(b))) is None:
+        name, n = f"buf_{n}", n+1
+        if not is_out: bufs_to_save[name] = b
+    bufs[key] = (name, size, bu.dtype, key)
+    return name
 
-def jit_model(model, *args) -> Tuple[TinyJit,Dict[int,str]]:
+  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))
+
+  return functions, statements, {name:(size, dtype, key) for name, size, dtype, key in bufs.values()}, bufs_to_save
+
+def jit_model(model, *args) -> Tuple[UOp, List[Buffer]]:
   assert hasattr(model, "forward") or callable(model), "model needs a forward function"
   @TinyJit
   def run(*x):
@@ -50,20 +52,10 @@ def jit_model(model, *args) -> Tuple[TinyJit,Dict[int,str]]:
     out = [out] if isinstance(out, Tensor) else out
     return [o.realize() for o in out]
 
-  # twice to run the JIT
+  # run twice to trigger JIT capture
   for _ in range(2): the_output = run(*args)
-  special_names = {}
-
-  # hack to put the inputs back
-  for (j,i),idx in run.input_replace.items():
-    realized_input = args[idx].uop.base.realized
-    run.jit_cache[j].bufs[i] = realized_input
-    special_names[id(realized_input)] = f'input{idx}'
-
-  # TODO: fetch this from the jit in self.input_replace and self.ret (hint: use get_parameters on self.ret)
-  for i, output in enumerate(the_output):
-    special_names[id(output.uop.base.realized)] = f'output{i}'
-  return run, special_names
+  assert run.captured is not None
+  return run.captured.linear, [o.uop.base.realized for o in the_output]
 
 def export_model_clang(functions:Dict[str,str], statements:Dict[str,Tuple[str,int,int]], bufs:Dict[str,Tuple[str,int,int]],
   bufs_to_save:Dict[str,Tensor], input_names:List[str], output_names:List[str], weight_names={}, model_name="model", symbolic_vars={}, wasm=False) -> str:
@@ -249,12 +241,12 @@ def export_model(model, target:str, *inputs, model_name: Optional[str] = "model"
   assert Device.DEFAULT in EXPORT_SUPPORTED_DEVICE, f"only {', '.join(EXPORT_SUPPORTED_DEVICE)} are supported"
 
   # NOTE: CPU_COUNT=1, since export does not support threading
-  with Context(JIT=2, CPU_COUNT=1): run,special_names = jit_model(model, *inputs)
-  functions, statements, bufs, bufs_to_save = compile_net(run, special_names)
+  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()}
-  input_names = [name for _,name in special_names.items() if "input" in name]
-  output_names = [name for _,name in special_names.items() if "output" in name]
+  input_names = [f"input{i}" for i in range(len(inputs))]
+  output_names = [f"output{i}" for i in range(len(output_bufs))]
 
   # handle symbolic variables; TODO: refactor to fix some of this stuff upstream in tinygrad
   symbolic_vars = OrderedDict()

extra/gemm/cdna_asm_gemm.py

@@ -6,7 +6,7 @@ from tinygrad.renderer import Estimates
 from tinygrad.helpers import getenv, all_same, DEBUG
 from tinygrad.runtime.support.compiler_amd import HIPCCCompiler
 from tinygrad.runtime.autogen.amd.cdna.ins import *
-from examples.mlperf.models.flat_llama import FP8_DTYPE, FP8_GRAD_DTYPE, matmul, quantize_fp8
+from examples.mlperf.models.flat_llama import FP8_DTYPE, FP8_GRAD_DTYPE, quantize_fp8
 
 # ** CDNA4 assembly gemm
 
@@ -2628,15 +2628,15 @@ def custom_asm_gemm(C:UOp, A:UOp, B:UOp, dname:str) -> UOp:
 # ** FP8 GEMM custom kernel
 
 @functools.cache
-def custom_hk_fp8_gemm(C:UOp, A:UOp, B:UOp, S:UOp, dname:str) -> UOp:
-  # A is (batch, M, K), B is (N, K) transposed, S is combined scale (scalar float)
+def custom_hk_fp8_gemm(C:UOp, A:UOp, B:UOp, X_s:UOp, W_s:UOp, dname:str) -> UOp:
+  # A is (batch, M, K), B is (N, K) transposed, X_s is x_scale, W_s is w_scale — kernel multiplies by both
   M, K = A.shape[0]*A.shape[1], A.shape[2]
   N, K2 = B.shape[(1 if B.ndim == 3 else 0):]
   assert K == K2, f"{A.shape} {B.shape}"
   block_size = 256
   threads = UOp.special(64 * 8, "lidx0")
   workgroups = UOp.special((M // block_size) * (N // block_size), "gidx0")
-  sink = UOp.sink(C.base, A.base, B.base, S.base, threads, workgroups,
+  sink = UOp.sink(C.base, A.base, B.base, X_s.base, W_s.base, threads, workgroups,
                   arg=KernelInfo(f"hk_fp8_gemm_{M}_{N}_{K}", estimates=Estimates(ops=2*M*N*K, mem=(M*K+N*K)*A.dtype.itemsize+M*N*C.dtype.itemsize)))
   kittens_path = pathlib.Path(__file__).parent.parent/"thunder"/"amd"
   src = (kittens_path/"gemm_fp8.cpp").read_text()
@@ -2698,19 +2698,19 @@ def custom_uop_gemm(C:UOp, A:UOp, B:UOp) -> UOp:
 
 def custom_gemm_bw(gradient:UOp, kernel:UOp):
   inputs = kernel.src[1:]
-  # fp8 scaled gemm has 4 inputs (out, a, b, scale), others have 3 (out, a, b)
-  if len(inputs) == 4:
-    out, a, b, scale = inputs
-    a_t, b_t, g_t, s_t = Tensor(a, device=a.device), Tensor(b, device=a.device), Tensor(gradient, device=a.device), Tensor(scale, device=a.device)
+  # fp8 scaled gemm has 5 inputs (out, a, b, x_scale, w_scale), others have 3 (out, a, b)
+  if len(inputs) == 5:
+    out, a, b, s_x, s_w = inputs
+    a_t, b_t, g_t = Tensor(a, device=a.device), Tensor(b, device=a.device), Tensor(gradient, device=a.device)
+    s_x_t, s_w_t = Tensor(s_x, device=a.device), Tensor(s_w, device=a.device)
     g_t = g_t[:a.shape[0]]
-    # backward GEMMs in fp8 with scale applied inside kernel to prevent bf16 overflow
     g_fp8, g_scale, _ = quantize_fp8(g_t)
-    bw_scale = g_scale * s_t
-    # dgrad: g_fp8 @ weight (asm_gemm computes a@b)
-    grad_a = asm_gemm(g_fp8, b_t, combined_scale=bw_scale)
-    # wgrad: g_fp8.T @ activation = (N, batch*seq) @ (batch*seq, K) → use permute to preserve sharding
-    grad_b = asm_gemm(g_fp8.permute(2, 0, 1).reshape(g_t.shape[-1], -1), a_t.reshape(-1, a_t.shape[-1]), combined_scale=bw_scale)
-    return (None, grad_a.uop, grad_b.uop, None)
+    # dgrad: uses g_scale * x_scale * w_scale
+    grad_a = asm_gemm(g_fp8, b_t, x_scale=g_scale * s_x_t, w_scale=s_w_t)
+    # wgrad: no w_scale
+    _one = Tensor(1.0, dtype=dtypes.float, device=a.device)
+    grad_b = asm_gemm(g_fp8.permute(2, 0, 1).reshape(g_t.shape[-1], -1), a_t.reshape(-1, a_t.shape[-1]), x_scale=g_scale * s_x_t, w_scale=_one)
+    return (None, grad_a.uop, grad_b.uop, None, None)
   else:
     out, a, b = inputs
     assert all_same([gradient.device, a.device, b.device, out.device])
@@ -2725,7 +2725,7 @@ def custom_gemm_bw(gradient:UOp, kernel:UOp):
 
 # ** main gemm function
 
-def asm_gemm(a:Tensor, b:Tensor, combined_scale:Tensor|None=None) -> Tensor:
+def asm_gemm(a:Tensor, b:Tensor, x_scale:Tensor|None=None, w_scale:Tensor|None=None) -> Tensor:
   assert can_use_asm_gemm(a, b), f"{counters['todos'][-1]}"
   counters["used"] += 1
   unfold_batch = a.ndim == 3 and isinstance(a.device, tuple) and a.uop.axis == 2 and b.uop.axis == 0
@@ -2757,10 +2757,12 @@ def asm_gemm(a:Tensor, b:Tensor, combined_scale:Tensor|None=None) -> Tensor:
   renderer = Device[dname:=(a.device[0] if is_multi else a.device)].renderer
   dname, arch = dname.split(":")[0], renderer.target.arch
   if arch.startswith("gfx950") and getenv("USE_ASM", 1):
-    # fp8 gemm computes [email protected], with optional combined scale applied inside kernel before bf16 store
+    # fp8 gemm computes [email protected], kernel multiplies output by x_scale * w_scale before bf16 store
     if a.dtype == FP8_DTYPE:
-      scale = combined_scale if combined_scale is not None else Tensor(1.0, dtype=dtypes.float, device=a.device)
-      out = Tensor.custom_kernel(out, a, b.T, scale, fxn=functools.partial(custom_hk_fp8_gemm, dname=dname), grad_fxn=custom_gemm_bw)[0]
+      _one = lambda: Tensor(1.0, dtype=dtypes.float, device=a.device)
+      xs = x_scale if x_scale is not None else _one()
+      ws = w_scale if w_scale is not None else _one()
+      out = Tensor.custom_kernel(out, a, b.T, xs, ws, fxn=functools.partial(custom_hk_fp8_gemm, dname=dname), grad_fxn=custom_gemm_bw)[0]
     else:
       out = Tensor.custom_kernel(out, a, b, fxn=functools.partial(custom_asm_gemm, dname=dname), grad_fxn=custom_gemm_bw)[0]
   else:

extra/gemm/tinygrad_nv_matmul.py

@@ -29,7 +29,7 @@ if __name__ == "__main__":
       Opt(op=OptOps.LOCAL, axis=0, amt=2),
     ]
   k.apply_opts(opts)
-  prg = get_program(k.ast, k.opts, k.applied_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)

extra/thunder/amd/gemm_fp8.cpp

@@ -93,7 +93,7 @@ constexpr int NUM_WARPS = 8;
 
 using G = kittens::group<NUM_WARPS>;
 
-__global__ __launch_bounds__(512, 2) void hk_fp8_gemm(bf16 *C_ptr, fp8e4m3 *A_ptr, fp8e4m3 *B_ptr, float *scale_ptr) {
+__global__ __launch_bounds__(512, 2) void hk_fp8_gemm(bf16 *C_ptr, fp8e4m3 *A_ptr, fp8e4m3 *B_ptr, float *x_scale_ptr, float *w_scale_ptr) {
     constexpr int M = GEMM_M, N = GEMM_N, K = GEMM_K;
 
     kittens::gl<fp8e4m3, 1, 1, M, K> A{A_ptr, nullptr, nullptr, nullptr, nullptr};
@@ -332,8 +332,8 @@ __global__ __launch_bounds__(512, 2) void hk_fp8_gemm(bf16 *C_ptr, fp8e4m3 *A_pt
         __builtin_amdgcn_s_barrier();
     }
 
-    // apply combined scale (x_scale * w_scale) before bf16 store to prevent overflow
-    float scale = *scale_ptr;
+    // apply x_scale * w_scale before bf16 store to prevent overflow
+    float scale = *x_scale_ptr * *w_scale_ptr;
     mul(cA, cA, scale);
     mul(cB, cB, scale);
     mul(cC, cC, scale);

extra/torch_backend/example.py

@@ -24,7 +24,7 @@ if __name__ == "__main__":
   kernel_count = GlobalCounters.kernel_count
   assert kernel_count > 0, "No kernels, test failed"
   # NOTE: this is 124 on torch 2.10.0
-  expected_kernels = 334
+  expected_kernels = 355
   expectation = f"ResNet18 kernels are {kernel_count} vs {expected_kernels} expected."
   if kernel_count < expected_kernels: warnings.warn(f"{expectation} Expectation can be lowered.", UserWarning)
   assert kernel_count <= expected_kernels, f"{expectation}"

extra/torch_backend/test_kernel_fusion.py

@@ -23,7 +23,7 @@ class TestKernelFusionRegression(unittest.TestCase):
     def fn():
       x = torch.randn(128, 128, device=device)
       return (x + 1.0) * 2.0 - 0.5
-    self._check_kernel_count(fn, 7)
+    self._check_kernel_count(fn, 6)
 
   def test_relu_fusion(self):
     def fn():
@@ -31,7 +31,7 @@ class TestKernelFusionRegression(unittest.TestCase):
       conv = torch.nn.Conv2d(3, 16, 3, padding=1).to(device)
       with torch.no_grad():
         return torch.nn.functional.relu(conv(x))
-    self._check_kernel_count(fn, 8)
+    self._check_kernel_count(fn, 7)
 
   def test_batchnorm_fusion(self):
     def fn():
@@ -41,26 +41,26 @@ class TestKernelFusionRegression(unittest.TestCase):
       bn.eval()
       with torch.no_grad():
         return torch.nn.functional.relu(bn(conv(x)))
-    self._check_kernel_count(fn, 12)
+    self._check_kernel_count(fn, 11)
 
   def test_reduce_fusion(self):
     def fn():
       x = torch.randn(64, 64, device=device)
       return (x * 2.0).sum()
-    self._check_kernel_count(fn, 7)
+    self._check_kernel_count(fn, 6)
 
   def test_matmul_elementwise_fusion(self):
     def fn():
       x = torch.randn(32, 32, device=device)
       w = torch.randn(32, 32, device=device)
       return torch.nn.functional.relu(x @ w + 1.0)
-    self._check_kernel_count(fn, 9)
+    self._check_kernel_count(fn, 8)
 
   def test_pooling_fusion(self):
     def fn():
       x = torch.randn(1, 8, 16, 16, device=device)
       return torch.nn.functional.max_pool2d(x * 2.0, 2)
-    self._check_kernel_count(fn, 7)
+    self._check_kernel_count(fn, 6)
 
   def test_residual_add_relu_fusion(self):
     def fn():
@@ -68,7 +68,7 @@ class TestKernelFusionRegression(unittest.TestCase):
       identity = torch.randn(1, 8, 16, 16, device=device)
       out = x + identity
       return torch.nn.functional.relu(out)
-    self._check_kernel_count(fn, 9)
+    self._check_kernel_count(fn, 8)
 
   def test_inplace_add_relu_fusion(self):
     def fn():
@@ -76,7 +76,7 @@ class TestKernelFusionRegression(unittest.TestCase):
       y = torch.randn(1, 16, 32, 32, device=device)
       x += y
       return torch.nn.functional.relu(x)
-    self._check_kernel_count(fn, 9)
+    self._check_kernel_count(fn, 8)
 
   def test_conv_bn_add_relu_fusion(self):
     def fn():
@@ -89,7 +89,7 @@ class TestKernelFusionRegression(unittest.TestCase):
         out = bn(conv(x))
         out += identity
         return torch.nn.functional.relu(out)
-    self._check_kernel_count(fn, 14)
+    self._check_kernel_count(fn, 13)
 
   def test_multiple_inplace_ops_fusion(self):
     def fn():
@@ -97,7 +97,7 @@ class TestKernelFusionRegression(unittest.TestCase):
       x += 1.0
       x *= 2.0
       return torch.nn.functional.relu(x)
-    self._check_kernel_count(fn, 6)
+    self._check_kernel_count(fn, 5)
 
   def test_view_inplace_no_fusion_break(self):
     def fn():
@@ -105,7 +105,7 @@ class TestKernelFusionRegression(unittest.TestCase):
       view = x[1:3]
       view += 1.0
       return x.sum()
-    self._check_kernel_count(fn, 10)
+    self._check_kernel_count(fn, 8)
 
   def test_batchnorm_running_stats_update(self):
     def fn():
@@ -114,7 +114,7 @@ class TestKernelFusionRegression(unittest.TestCase):
       bn.train()
       with torch.no_grad():
         return bn(x)
-    self._check_kernel_count(fn, 10)
+    self._check_kernel_count(fn, 9)
 
   # this is a minimal extra/other_mnist/beautiful_mnist_torch.py to cover fusion for training with optimizer
   def test_mnist_training_fusion(self):
@@ -135,7 +135,7 @@ class TestKernelFusionRegression(unittest.TestCase):
       loss.backward()
       optimizer.step()
       return loss
-    self._check_kernel_count(fn, 26)
+    self._check_kernel_count(fn, 25)
 
 if __name__ == "__main__":
   unittest.main()

extra/viz/README.md

@@ -1,49 +0,0 @@
-A command line tool for exploring the VIZ trace.
-
-# Lightweight tracing
-
-Supported on all backends.
-
-Flags: VIZ=-1 to only save the trace to a file, VIZ=1 also launches a web server.
-
-1. Set VIZ to -1 to save the trace.
-2. Use `extra/viz/cli.py` to inspect the trace files.
-
-## Inspect runtime profiling
-
-Use `extra/viz/cli.py --profile` to list all sources.
-
-List top slowest kernels on a source: `--profile -s "AMD"`
-List samples of a kernel on a source: `--profile -s "AMD" -i E_3 | head 4`
-
-## Inspect codegen and PatternMatcher
-
-Use `extra/viz/cli.py --rewrites` to list all sources.
-
-List all codegen steps for a kernel: `--rewrites -s E_3`
-Get source code: `--rewrites -s E_3 -i "View Source"`
-Inspect a graph rewrite: `--rewrites -s E_3 -i "initial symbolic"`
-
-## SQTT tracing
-
-Supported on AMD for RDNA3 and RDNA4 (best) and CDNA (developing).
-
-Flags: VIZ=-2 to save SQTT trace to a file. VIZ=2 also launches a web server. View other flags in tinygrad/runtime/ops_amd.py to configure SQTT as needed.
-
-Use `extra/viz/cli.py --profile | grep SQTT` to view all available SQTT traces.
-You can select a specific trace with --source, Example workflow:
-
-```bash
-# Run amd_asm_matmul with VIZ=-2 to capture the trace
-VIZ=-2 python extra/gemm/amd_asm_matmul.py
-
-# View barriers
-extra/viz/cli.py --profile -s "kernel SQTT SE:0 PKTS" | rg BARRIER | head -10
-
-# Get bank conflicts from performance counters
-
-python extra/viz/cli.py -p -s "kernel PMC" -i "SQC_LDS_BANK_CONFLICT"
-
-# Find the EXEC corresponding to a DISPATCH at cycle 410
-extra/viz/cli.py --profile -s "kernel SQTT SE:0 PKTS" | awk '/EXEC/ && $1 - $5 == 410'
-```

extra/viz/cli.py

@@ -1,186 +0,0 @@
-#!/usr/bin/env python3
-import argparse, pathlib, signal, sys, struct, json, itertools
-if hasattr(signal, "SIGPIPE"): signal.signal(signal.SIGPIPE, signal.SIG_DFL)
-from typing import Iterator
-from tinygrad.viz import serve as viz
-from tinygrad.uop.ops import RewriteTrace
-from tinygrad.helpers import temp, ansistrip, colored, time_to_str, ansilen, ProfilePointEvent, ProfileRangeEvent, TracingKey, unwrap
-
-# profile decoder used in CLI and tests
-def decode_profile(data:bytes) -> dict:
-  ret, off = data, 0
-  def u(fmt:str) -> tuple:
-    nonlocal off
-    vals = struct.unpack_from(fmt, ret, off)
-    off += struct.calcsize(fmt)
-    return vals
-  total_dur, global_peak, index_len, layout_len = u("<IQII")
-  strings, dtypes, markers = json.loads(ret[off:off+index_len]).values()
-  off += index_len
-  layout:dict[str, dict] = {}
-  # 0 means None, otherwise it's an enum value
-  def option(i:int) -> int|None: return None if i == 0 else i-1
-  for _ in range(layout_len):
-    klen = u("<B")[0]
-    k = ret[off:off+klen].decode()
-    off += klen
-    v:dict = {"events":[]}
-    layout[k] = v
-    event_type, event_count = u("<BI")
-    if event_type == 0:
-      for _ in range(event_count):
-        name, ref, key, st, dur, fmt = u("<IIIIfI")
-        v["events"].append({"name":strings[name], "ref":option(ref), "key":option(key), "st":st, "dur":dur, "fmt":strings[fmt]})
-    else:
-      v["linear"] = u("<B")[0]
-      v["peak"] = u("<Q")[0]
-      for _ in range(event_count):
-        if v["linear"]:
-          ts, value = u("<IQ")
-          v["events"].append({"event":"freq", "ts":ts, "value":value})
-        else:
-          alloc, ts, key = u("<BII")
-          if alloc: v["events"].append({"event":"alloc", "ts":ts, "key":key, "arg": {"dtype":strings[u("<I")[0]], "sz":u("<Q")[0]}})
-          else: v["events"].append({"event":"free", "ts":ts, "key":key, "arg": {"users":[u("<IIIB") for _ in range(u("<I")[0])]}})
-  return {"dur":total_dur, "peak":global_peak, "layout":layout, "markers":markers}
-
-def get(data:dict, key:str):
-  for k,v in data.items():
-    if ansistrip(k) == key: return v
-  import difflib
-  match = difflib.get_close_matches(key, [ansistrip(k) for k in data], n=1, cutoff=0.6)
-  raise RuntimeError(f'item "{key}" not found in list'+(f", did you mean {match[0]!r}?" if match else ''))
-
-def main(args) -> None:
-  viz.load_rewrites(viz_data:=viz.VizData(viz.load_pickle(args.rewrites_path, default=RewriteTrace([], [], {}))))
-
-  def format_colored(s:str) -> str: return ansistrip(s) if args.no_color else s
-
-  if args.profile:
-    events:list = viz.load_pickle(args.profile_path, default=[])
-    if (profile_bytes:=viz.get_profile(viz_data, events)) is None: raise RuntimeError(f"empty profile in {args.profile_path}")
-    profile = decode_profile(profile_bytes)
-    profile["layout"].update([(f'{c["name"][5:]}{" SQTT" if s["name"].endswith("PKTS") else ""} {s["name"]}', s["data"]) for c in viz_data.ctxs
-                              if c["name"].startswith("SQTT") for s in c["steps"] if s["name"].endswith(("PMC", "PKTS"))])
-    if args.src is None:
-      for k in profile["layout"]:
-        print(f"  {format_colored(k)}")
-      return None
-
-    # ** SQTT printer
-    data = get(profile["layout"], args.src)
-    if "SQTT" in args.src:
-      # modern terminals support 24-bit color
-      def hex_colored(st:str, color:str) -> str: return f"\x1b[38;2;{int(color[1:3],16)};{int(color[3:5],16)};{int(color[5:7],16)}m{st}\x1b[0m"
-      print(f"{'Clk':<12} {'Unit':<20} {'Op':<22} {'Dur':<4} {'Delay':<4} {'Info'}")
-      print("-" * 100)
-      pc_map:dict[int, str] = {}
-      pkt_idxs:dict[str, itertools.count] = {}
-      dispatch_to_inst:dict[str, tuple[str, int]] = {}
-      inst_st:int|None = None
-      for e in viz.sqtt_timeline(*data):
-        if isinstance(e, ProfilePointEvent) and e.key == 'pcMap': pc_map = e.arg
-        if not isinstance(e, ProfileRangeEvent): continue
-        if inst_st is None: inst_st = int(e.st)
-        assert isinstance(e.name, TracingKey)
-        op_name, info = e.name.display_name, e.name.ret or ""
-        color = next((v for k,v in viz.wave_colors.items() if k in op_name), None)
-        op_str = hex_colored(op_name, color) if color and not args.no_color else op_name
-        phase, delay = None, 0
-        idx = next(pkt_idxs.setdefault(e.device, itertools.count()))
-        if e.device.startswith("WAVE"):
-          inst = f"0x{(pc:=int(info.replace('PC:', ''))):05x} {pc_map[pc]}" if info else f"{'':7} {op_name}"
-          dispatch_to_inst[f"{e.device}-{idx}"] = (inst, int(e.st))
-          phase = "DISPATCH"
-        if info.startswith("LINK:"):
-          inst, dispatch_st = dispatch_to_inst[info.replace("LINK:", "")]
-          phase, delay = "EXEC", int(e.st) - dispatch_st
-        if inst and phase: info = f"{phase:<8} {inst}"
-        unit = e.device.replace(" ", "-")
-        print(f"{int(e.st)-inst_st:<12} {unit:<20} {op_str}{' '*(22-ansilen(op_str))} {int(unwrap(e.en)-e.st):<4} {str(delay or ''):<4} {info}")
-      return None
-
-    # ** PMC printer
-    if "PMC" in args.src:
-      pmc = viz.unpack_pmc(data)
-      cols = pmc["cols"]
-      rows:list = []
-      for r in pmc["rows"]:
-        if args.item is None: rows.append(r[:2])
-        elif args.item == r[0]:
-          rows = r[2]["rows"] if len(r) > 2 else [r[:2]]
-          cols = r[2]["cols"] if len(r) > 2 else cols
-      from tabulate import tabulate
-      print(tabulate(rows, headers=cols, tablefmt="github"))
-      return None
-
-    # ** Profiler printer
-    agg:dict[str, tuple[float, int]] = {}
-    total = 0
-    for e in data.get("events", []):
-      et = e["dur"] * 1e-6
-      if args.item is not None:
-        if ansistrip(e["name"]) == args.item:
-          ptm = colored(time_to_str(et, w=9), "yellow" if et > 0.01 else None)
-          name = e["name"] + (" " * (46 - ansilen(e["name"])))
-          print(f"{format_colored(name)} {ptm}/{et*1e3:9.2f}ms  " + e.get("fmt", "").replace("\n", " | ") + "  ")
-      else:
-        t, c = agg.get(e["name"], (0.0, 0))
-        agg[e["name"]] = (t+et, c+1)
-        total += et
-    if agg and total > 0:
-      from tabulate import tabulate
-      items = sorted(agg.items(), key=lambda kv:kv[1][0], reverse=True)
-      num_rows = 20
-      table = [[format_colored(name), time_to_str(t, w=9), c, f"{(t/total*100.0):.2f}%"] for name,(t,c) in items[:num_rows]]
-      if items[num_rows:]:
-        other_t = sum(t for _,(t,_) in items[num_rows:])
-        other_c = sum(c for _,(_,c) in items[num_rows:])
-        table.append(["Other", time_to_str(other_t, w=9), other_c, f"{(other_t/total*100.0):.2f}%"])
-      print(tabulate(table, headers=["name", "total", "count", "pct"], tablefmt="github"))
-    return None
-
-  # ** Graph rewrites printer
-  rewrites = {c["name"]:{s["name"]:s for s in c["steps"]} for c in viz_data.ctxs if c.get("steps")}
-  if args.src is None:
-    for k in rewrites: print(f"  {format_colored(k)}")
-    return None
-  steps = get(rewrites, args.src)
-  if args.item is None:
-    for k,v in steps.items(): print(" "*v["depth"]+k+(f" - {v['match_count']}" if v.get('match_count', 0) else ''))
-  else:
-    data = viz.get_render(data, get(steps, args.item)["query"])
-    if isinstance(data.get("value"), Iterator):
-      for m in data["value"]:
-        if m.get("uop"): print(f"Input UOp:\n{m['uop']}")
-        if m.get("diff"):
-          loc = pathlib.Path(m["upat"][0][0])
-          print(f"Rewrite at {loc.parent.name}/{loc.name}:{m['upat'][0][1]}\n{m['upat'][1]}")
-          for line in m["diff"]:
-            print(line if args.no_color else colored(line, "red" if line.startswith("-") else "green" if line.startswith("+") else None))
-    if data.get("src") is not None: print(data["src"])
-
-def get_arg_parser() -> argparse.ArgumentParser:
-  parser = argparse.ArgumentParser(add_help=False)
-  g_mode = parser.add_argument_group("mode")
-  g_mode.add_argument("-p", "--profile", action="store_true", help="View profile")
-  g_mode.add_argument("-r", "--rewrites", action="store_true", help="View graph rewrites")
-  g_opts = parser.add_argument_group("optional args")
-  g_opts.add_argument("-s", "--src", type=str, default=None, metavar="NAME", help="Select a data source (default: list all sources)")
-  g_opts.add_argument("-i", "--item", type=str, default=None, metavar="NAME", help="Select an item within the source (default: list all items)")
-  g_opts.add_argument("--no-color", action="store_true", help="Turn off colored names")
-  g_opts.add_argument("--profile-path", type=pathlib.Path, metavar="PATH", help="Path to profile.pkl (optional file, default: latest profile)",
-                      default=pathlib.Path(temp("profile.pkl", append_user=True)))
-  g_opts.add_argument("--rewrites-path", type=pathlib.Path, metavar="PATH", help="Path to rewrites.pkl (optional file, default: latest rewrites)",
-                      default=pathlib.Path(temp("rewrites.pkl", append_user=True)))
-  g_opts.add_argument("-h", "--help", action="help", help="show this help message and exit")
-  return parser
-
-if __name__ == "__main__":
-  args = get_arg_parser().parse_args()
-  if not args.profile and not args.rewrites:
-    get_arg_parser().print_help()
-    sys.exit(0)
-
-  try: main(args)
-  except KeyboardInterrupt: pass

pyproject.toml

@@ -19,11 +19,11 @@ build-backend = "setuptools.build_meta"
 include-package-data = true
 packages = [
   'tinygrad',
-  'tinygrad.apps',
   'tinygrad.codegen',
   'tinygrad.codegen.opt',
   'tinygrad.codegen.late',
   'tinygrad.engine',
+  'tinygrad.llm',
   'tinygrad.mixin',
   'tinygrad.nn',
   'tinygrad.renderer',
@@ -38,6 +38,7 @@ packages = [
   'tinygrad.runtime.graph',
   'tinygrad.runtime.support',
   'tinygrad.runtime.support.am',
+  'tinygrad.runtime.support.mlx',
   'tinygrad.runtime.support.nv',
   'tinygrad.schedule',
   'tinygrad.uop',
@@ -74,7 +75,7 @@ testing_minimal = [
   "hypothesis>=6.148.9",
   "z3-solver<4.15.4",  # 4.15.4 has a segfault when creating many z3.Context()
 ]
-testing_unit = ["tinygrad[testing_minimal]", "tqdm", "safetensors", "tabulate", "openai", "gguf>=0.18"]
+testing_unit = ["tinygrad[testing_minimal]", "tqdm", "safetensors", "tabulate", "openai", "gguf>=0.18", "capstone"]
 testing = [
   "tinygrad[testing_unit]",
   "pillow",
@@ -92,7 +93,6 @@ testing = [
   "networkx",
   "nibabel",
   "bottle",
-  "capstone",
   "pycocotools",
   "boto3",
   "pandas",
@@ -112,9 +112,9 @@ docs = [
 [tool.mutmut]
 paths_to_mutate = ["tinygrad/"]
 do_not_mutate = [
-  "tinygrad/apps/*",
   "tinygrad/codegen/*",
   "tinygrad/engine/*",
+  "tinygrad/llm/*",
   "tinygrad/nn/*",
   "tinygrad/renderer/*",
   "tinygrad/runtime/*",

sz.py

@@ -56,7 +56,7 @@ def gen_diff(table_old, table_new):
 
 def display_diff(diff): return "+"+str(diff) if diff > 0 else str(diff)
 
-NONCORE_DIRS = {"tinygrad/apps", "tinygrad/nn", "tinygrad/renderer", "tinygrad/runtime", "tinygrad/viz"}
+NONCORE_DIRS = {"tinygrad/llm", "tinygrad/nn", "tinygrad/renderer", "tinygrad/runtime", "tinygrad/viz"}
 
 if __name__ == "__main__":
   if len(sys.argv) == 3:

test/amd/test_mockgpu_invalid.py

@@ -37,8 +37,7 @@ dev.synchronize()
 '''
 
     env = os.environ.copy()
-    env["AMD"] = "1"
-    env["MOCKGPU"] = "1"
+    env["DEV"] = "MOCKKFD+AMD"
     env["HCQDEV_WAIT_TIMEOUT_MS"] = "10000"
 
     st = time.perf_counter()

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=AMD MOCKGPU=1 python -m pytest test/amd/test_sqtt_encoder.py -v
+Run with: DEV=MOCKKFD+AMD python -m pytest test/amd/test_sqtt_encoder.py -v
 """
 import ctypes, unittest
 from tinygrad.helpers import Context

test/amd/test_sqttmap.py

@@ -1,5 +1,5 @@
 # test to compare every packet with the rocprof decoder
-import unittest, pickle, contextlib, io
+import unittest, pickle
 from typing import Iterator
 from pathlib import Path
 from tinygrad.helpers import DEBUG, getenv, temp, ansistrip
@@ -7,17 +7,11 @@ 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 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 run_cli(*cli_args) -> str:
-  from extra.viz.cli import main, get_arg_parser
-  args = get_arg_parser().parse_args(cli_args)
-  with contextlib.redirect_stdout(buf:=io.StringIO()):
-    main(args)
-  return buf.getvalue().strip()
-
 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

test/backend/test_asm_gemm.py

@@ -1,7 +1,7 @@
 import unittest
 from tinygrad import Tensor, Device, dtypes, Context
 from tinygrad.device import is_dtype_supported
-from tinygrad.helpers import getenv, system
+from tinygrad.helpers import getenv, system, DEV
 from extra.gemm.cdna_asm_gemm import asm_gemm
 from test.helpers import needs_second_gpu
 from examples.mlperf.models.flat_llama import FP8_DTYPE
@@ -131,7 +131,7 @@ class TestGemmLlama(unittest.TestCase):
   dtype = dtypes.bfloat16
 
   def setUp(self):
-    if not is_cdna4() or getenv("MOCKGPU"):
+    if not is_cdna4() or DEV.interface.startswith("MOCK"):
       self.skipTest("very slow on non mi350x")
 
   def test_empty(self): asm_gemm(Tensor.empty(N:=getenv("N", 4096), N, dtype=self.dtype), Tensor.empty(N, N, dtype=self.dtype)).realize()

test/backend/test_const_folding.py

@@ -1,7 +1,7 @@
 import unittest, math
 from tinygrad import Tensor, Device, dtypes
 from tinygrad.dtype import DTYPES_DICT
-from tinygrad.uop.ops import Ops
+from tinygrad.uop.ops import Ops, UOp
 from tinygrad.device import is_dtype_supported
 import numpy as np
 from test.helpers import not_support_multi_device
@@ -163,6 +163,11 @@ class TestMultiConstFolding(unittest.TestCase):
     np.testing.assert_equal((t ** one).numpy(), np.arange(16))
     np.testing.assert_equal((one ** t).numpy(), [1] * 16)
 
+class TestThreefryConstFolding(unittest.TestCase):
+  def test_threefry(self):
+    x = UOp.const(dtypes.uint64, 5, Device.DEFAULT, ()).threefry(UOp.const(dtypes.uint64, 10, Device.DEFAULT, ()))
+    self.assertIs(x.simplify().op, Ops.CONST)
+
 class TestTautologicalCompare(unittest.TestCase):
   # without const folding, these would have triggered -Wtautological-compare in clang
   def test_lt_false(self):

test/backend/test_custom_kernel.py

@@ -12,7 +12,7 @@ def custom_arange_kernel(C:UOp) -> UOp:
 def custom_eye_kernel(C:UOp) -> UOp:
   i = UOp.range(C.shape[0], 0)
   j = UOp.range(C.shape[1], 1)
-  return C[i, j].store((i.eq(j)).cast(C.dtype.base)).end(i, j).sink(arg=KernelInfo(name=f"custom_eye_{C.size}"))
+  return C[i, j].store((i.eq(j)).cast(C.dtype.base)).end(i, j).sink(arg=KernelInfo(name=f"custom_eye_{C.numel()}"))
 
 def custom_add_one_kernel(B:UOp, A:UOp) -> UOp:
   A,B = A.flatten(), B.flatten()
@@ -52,7 +52,7 @@ def flip_contract_kernel(dest:UOp, src:UOp):
   j = UOp.range(dest.shape[1], 1, AxisType.UPCAST)
   vec = src[i, j].contract(j)
   store = UOp.group(*[dest[i, k].store(vec.gep(3-k)) for k in range(4)])
-  return store.end(i, j).sink(arg=KernelInfo(name=f"flip_contract_{dest.size}", opts_to_apply=()))
+  return store.end(i, j).sink(arg=KernelInfo(name=f"flip_contract_{dest.numel()}", opts_to_apply=()))
 
 def slice_sum_kernel(dest:UOp, src:UOp):
   G = UOp.range(src.shape[0], 0)

test/backend/test_dtype_alu.py

@@ -1,7 +1,7 @@
 import unittest, operator, math
 from tinygrad import Context, Tensor, dtypes, Device
 from tinygrad.dtype import DType, truncate, fp8_to_float
-from tinygrad.helpers import CI, EMULATED_DTYPES, getenv
+from tinygrad.helpers import CI, EMULATED_DTYPES, DEV, getenv
 from tinygrad.tensor import _to_np_dtype
 from tinygrad.device import is_dtype_supported
 from tinygrad.runtime.ops_python import from_storage_scalar
@@ -32,7 +32,8 @@ unary_operations = [(Tensor.exp, np.exp), (Tensor.log, np.log), (Tensor.sin, np.
 #binary_operations.append(operator.truediv)
 
 # TODO: CI CUDA segfaults on sin, WEBGPU and NIR sines are not precise enough for large numbers
-if (getenv("MOCKGPU") and Device.DEFAULT in {"NV", "CUDA"}) or Device.DEFAULT == "WEBGPU" or isinstance(Device[Device.DEFAULT].renderer, NIRRenderer):
+if ((DEV.interface.startswith("MOCK") and Device.DEFAULT in {"NV", "CUDA"})
+    or Device.DEFAULT == "WEBGPU" or isinstance(Device[Device.DEFAULT].renderer, NIRRenderer)):
   unary_operations.remove((Tensor.sin, np.sin))
   unary_operations.remove((Tensor.cos, np.cos))
 

test/backend/test_edgecases.py

@@ -27,10 +27,10 @@ import numpy as np
 import torch
 from tinygrad import Tensor, dtypes, nn
 from tinygrad.device import Device
-from tinygrad.helpers import getenv
+from tinygrad.helpers import DEV
 from tinygrad.renderer.nir import NIRRenderer
 
-MOCKGPU = getenv("MOCKGPU")
+MOCKGPU = DEV.interface.startswith("MOCK")
 
 class TestNaNEdgeCases(unittest.TestCase):
   # we don't need more of these. it's unclear if torch's behavior is desired here

test/backend/test_interop.py

@@ -3,12 +3,12 @@ import unittest
 import torch
 import numpy as np
 
-from tinygrad.helpers import getenv, CI
+from tinygrad.helpers import CI, DEV
 from tinygrad.tensor import Tensor
 from tinygrad.device import Device
 from tinygrad.dtype import _from_torch_dtype, _to_torch_dtype
 
-MOCKGPU = getenv("MOCKGPU")
+MOCKGPU = DEV.interface.startswith("MOCK")
 
 @unittest.skipIf(Device.DEFAULT not in ["METAL", "CUDA"] or MOCKGPU, f"no support on {Device.DEFAULT}")
 class TestInterop(unittest.TestCase):

test/backend/test_jit.py

@@ -1,15 +1,15 @@
 #!/usr/bin/env python
-import unittest, functools
+import unittest
 import numpy as np
 
 from hypothesis import given, settings, strategies as strat
-from test.helpers import assert_jit_cache_len, not_support_multi_device, needs_second_gpu
+from test.helpers import assert_jit_cache_len, call_is_graph, not_support_multi_device, needs_second_gpu
 from tinygrad.tensor import Tensor
-from tinygrad.engine.jit import TinyJit, JitError, GraphRunner, MultiGraphRunner, graph_class
-from tinygrad.engine.realize import CompiledRunner, BufferCopy, BufferXfer
+from tinygrad.engine.jit import TinyJit, JitError, graph_class
 from tinygrad.device import Device
-from tinygrad.helpers import Context, JIT, GlobalCounters, getenv
+from tinygrad.helpers import Context, JIT, DEV, GlobalCounters
 from tinygrad.dtype import dtypes
+from tinygrad.uop.ops import Ops
 from extra.models.unet import ResBlock
 
 def _simple_test(add, extract=lambda x: x, N=10):
@@ -419,10 +419,10 @@ class TestJit(unittest.TestCase):
       if prev is not None: np.testing.assert_allclose(o, prev, atol=1e-4, rtol=1e-5)
       prev = o
 
-    graph_t = Device[Device.DEFAULT].graph.func if isinstance(Device[Device.DEFAULT].graph, functools.partial) else Device[Device.DEFAULT].graph
     # Checking that 2 graphs are inited.
-    assert isinstance(jf.jit_cache[0].prg, graph_t)
-    assert isinstance(jf.jit_cache[1].prg, graph_t)
+    assert len(jf.captured.linear.src) == 2
+    for si in jf.captured.linear.src:
+      assert call_is_graph(si)
 
   def test_jitted_clone(self):
     def f(a): return a.clone().realize()
@@ -583,7 +583,7 @@ class TestJitPrune(unittest.TestCase):
       a = Tensor.rand(16).realize()
       out = w2_prune(a)
       np.testing.assert_allclose(out.tolist(), [x*2+y for x,y in zip(weights.tolist(), a.tolist())])
-    assert len(w2_prune.captured.jit_cache) == 1
+    assert_jit_cache_len(w2_prune, 1)
 
   def test_prune_w_copy_correct(self):
     weights = Tensor.rand(16).realize()
@@ -617,7 +617,7 @@ class TestJitPrune(unittest.TestCase):
       out = w2_prune(a)
       np.testing.assert_allclose(out.tolist(), [x*2+y for x,y in zip(weights.tolist(), a.tolist())])
 
-    assert len(w2_prune.captured.jit_cache) == 1, "prune should have removed the copy"
+    assert_jit_cache_len(w2_prune, 1)
 
 class TestJitFree(unittest.TestCase):
   def test_free_intermediates(self):
@@ -688,8 +688,9 @@ class TestJitGraphSplit(unittest.TestCase):
     graph_t = graph_class(dev)
     if graph_t is None: return
 
-    got = f.jit_cache
+    got = f.captured.linear.src
     from tinygrad.runtime.graph.hcq import HCQGraph
+    from tinygrad.engine.jit import MultiGraphRunner
     if graph_t is HCQGraph:
       validate = hcqgraph
     elif issubclass(graph_t, MultiGraphRunner):
@@ -698,16 +699,16 @@ class TestJitGraphSplit(unittest.TestCase):
       validate = graph
 
     assert len(got) == len(validate), f"Expected {len(validate)} operations, got {len(got)}"
-    for expected, got in zip(validate, got):
+    for expected, si in zip(validate, got):
+      ast = si.src[0]
       if expected["type"] == "graph":
-        assert isinstance(got.prg, GraphRunner), f"Expected GraphRunner, got {type(got.prg)}"
-        assert len(got.prg.jit_cache) == expected["cnt"], f"Expected {expected['cnt']} operations in graph, got {len(got.prg.jit_cache)}"
+        assert call_is_graph(si), f"Expected graph, got {ast.op}"
+        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 isinstance(got.prg, CompiledRunner), f"Expected CompiledRunner, got {type(got.prg)}"
-      elif expected["type"] == "copy":
-        assert isinstance(got.prg, BufferCopy), f"Expected BufferCopy, got {type(got.prg)}"
-      elif expected["type"] == "xfer":
-        assert isinstance(got.prg, BufferXfer), f"Expected BufferXfer, got {type(got.prg)}"
+        assert ast.op in (Ops.SINK, Ops.PROGRAM, Ops.BEAM), f"Expected kernel, got {ast.op}"
+      elif expected["type"] in ("copy", "xfer"):
+        assert ast.op is Ops.COPY, f"Expected COPY, got {ast.op}"
 
   def ji_graph(self, cnt): return {"type": "graph", "cnt": cnt}
   def ji_comp(self): return {"type": "comp"}
@@ -812,7 +813,7 @@ class TestJitGraphSplit(unittest.TestCase):
       hcqgraph=[self.ji_graph(6)])
 
   @unittest.skip("this fails if you don't have SDMA or are using AMD_DISABLE_SDMA=1")
-  @unittest.skipIf(getenv("MOCKGPU"), "MockGPU does not support parallel copies")
+  @unittest.skipIf(DEV.interface.startswith("MOCK"), "MockGPU does not support parallel copies")
   def test_jit_multidev_copy(self):
     if Device.DEFAULT in {"CPU"}: raise unittest.SkipTest("CPU/LLVM is not a valid default device for this test (zero-copies)")
 

test/backend/test_linearizer.py

@@ -7,11 +7,12 @@ 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.helpers import Context, flatten, dedup, TC_SELECT, TC_OPT, getenv
+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
 from tinygrad.renderer.cstyle import CUDARenderer
-MOCKGPU = getenv("MOCKGPU")
+from test.helpers import replace_opts
+MOCKGPU = DEV.interface.startswith("MOCK")
 
 from tinygrad.uop.ops import print_uops # noqa: F401 # pylint: disable=unused-import
 
@@ -44,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(ast, renderer=Device[Device.DEFAULT].renderer, opts=[]).uops
+    uops = get_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).uops
     self.assertEqual(len([x for x in uops if x.op is Ops.CAST]), 1)
 
   @unittest.expectedFailure
@@ -52,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(ast, renderer=Device[Device.DEFAULT].renderer, opts=[]).uops
+    uops = get_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).uops
     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")
@@ -62,7 +63,7 @@ class TestLinearizer(unittest.TestCase):
     b = Tensor.empty(16)
     out = img.conv2d(w, b)
     ast = helper_linearizer_opt(out)
-    uops = get_program(ast, renderer=Device[Device.DEFAULT].renderer, opts=[]).uops
+    uops = get_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).uops
     # 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
@@ -83,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(ast, renderer=Device[Device.DEFAULT].renderer, opts=[]).uops
+    uops = get_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).uops
     ranges = [i for i,u in enumerate(uops) if u.op is Ops.RANGE]
     assert len(ranges) == 1 # NOTE: it collapses now
 
@@ -91,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(ast, renderer=Device[Device.DEFAULT].renderer, opts=[]).uops
+    uops = get_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).uops
     ranges = [i for i,u in enumerate(uops) if u.op is Ops.RANGE]
     assert len(ranges) == 1 # NOTE: it collapses now
 
@@ -99,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(ast, renderer=Device[Device.DEFAULT].renderer, opts=[]).uops
+    uops = get_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).uops
     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]])
@@ -112,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(ast, renderer=Device[Device.DEFAULT].renderer, opts=[]).uops
+    uops = get_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).uops
     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
@@ -122,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(ast, renderer=Device[Device.DEFAULT].renderer, opts=[]).uops
+    uops = get_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).uops
     ranges = [i for i,u in enumerate(uops) if u.op is Ops.RANGE]
     assert len(ranges) == 1 # NOTE: it collapses now
 
@@ -133,7 +134,7 @@ class TestLinearizer(unittest.TestCase):
     # these are of size 3 to avoid float4 coalesce
     r = a[:-1] + a[1:]
 
-    uops = get_program(r.schedule()[-1].ast, renderer=Device[Device.DEFAULT].renderer, opts=[Opt(op=OptOps.UPCAST, axis=0, arg=0)]).uops
+    uops = get_program(replace_opts(r.schedule()[-1].ast, [Opt(op=OptOps.UPCAST, axis=0, arg=0)]), renderer=Device[Device.DEFAULT].renderer).uops
     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?"
@@ -145,7 +146,7 @@ class TestLinearizer(unittest.TestCase):
     a, b = Tensor.randn(1).realize(), Tensor.randn(1).realize()
     r = a.expand([2]) + b.expand([2])
 
-    uops = get_program(r.schedule()[-1].ast, renderer=Device[Device.DEFAULT].renderer, opts=[Opt(op=OptOps.UPCAST, axis=0, arg=0)]).uops
+    uops = get_program(replace_opts(r.schedule()[-1].ast, [Opt(op=OptOps.UPCAST, axis=0, arg=0)]), renderer=Device[Device.DEFAULT].renderer).uops
     num_ops = len([uop for uop in uops if uop.op in GroupOp.ALU])
     assert num_ops <= 1, "more alu uops than needed"
 
@@ -154,8 +155,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(r.schedule()[-1].ast, renderer=Device[Device.DEFAULT].renderer,
-                       opts=[Opt(op=OptOps.UPCAST, axis=0, arg=0), Opt(op=OptOps.UNROLL, axis=0, arg=0)]).uops
+    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
     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
@@ -167,7 +168,7 @@ class TestLinearizer(unittest.TestCase):
   @unittest.skipUnless(Device.DEFAULT == "CPU", "test only for CPU")
   def test_upcast_with_locals_cpu(self):
     out = Tensor.ones(64,64).contiguous() @ Tensor.ones(64,64).contiguous()
-    prg = get_program(out.schedule()[-1].ast, opts=[Opt(OptOps.LOCAL, axis=0, arg=4)]).uops
+    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)
 
   @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals")
@@ -178,7 +179,7 @@ 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(r.schedule()[-1].ast, renderer=Device[Device.DEFAULT].renderer, opts=opts_to_apply)
+    program = get_program(replace_opts(r.schedule()[-1].ast, 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]
 
@@ -192,7 +193,7 @@ 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(r.schedule()[-1].ast, renderer=Device[Device.DEFAULT].renderer, opts=[Opt(op=OptOps.UPCAST, axis=0, arg=0)]).uops
+    uops = get_program(replace_opts(r.schedule()[-1].ast, [Opt(op=OptOps.UPCAST, axis=0, arg=0)]), renderer=Device[Device.DEFAULT].renderer).uops
     num_ops = len([uop for uop in uops if uop.op in GroupOp.ALU])
     assert num_ops == 0, "more alu uops than needed"
 
@@ -202,14 +203,14 @@ class TestLinearizer(unittest.TestCase):
       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(realized_ast, renderer=Device[Device.DEFAULT].renderer, opts=[])
+        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]
         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(realized_ast, renderer=Device[Device.DEFAULT].renderer, opts=[])
+      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]
       self.assertEqual(local[0].dtype.base, expected_dtype)
 
@@ -237,7 +238,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(ast, renderer=Device[Device.DEFAULT].renderer, opts=opt).uops
+    uops = get_program(replace_opts(ast, opt), renderer=Device[Device.DEFAULT].renderer).uops
     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)
@@ -257,7 +258,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(ast, renderer=Device[Device.DEFAULT].renderer, opts=[]).uops
+    uops = get_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).uops
     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]
@@ -290,7 +291,7 @@ class TestLinearizer(unittest.TestCase):
     sched_copy = sched[:]
     run_schedule(sched)
     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(sched_copy[-1].ast, renderer=Device[Device.DEFAULT].renderer, opts=())
+    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"
 
   def test_phi_simplification(self):
@@ -352,7 +353,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(ast, renderer=Device[Device.DEFAULT].renderer, opts=opt).uops
+    uops = get_program(replace_opts(ast, opt), renderer=Device[Device.DEFAULT].renderer).uops
     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]
@@ -420,7 +421,9 @@ def _helper_linearizer_opt_ast(realized_ast:UOp, real_bufs:list[Buffer], opts=[]
   device = real_bufs[0].device
   wanna_output = [np.array(x).flatten() for x in wanna_output]
 
-  def get_prg(opts): return CompiledRunner(replace(get_program(realized_ast, renderer=Device[Device.DEFAULT].renderer, opts=opts), device=device))
+  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))
 
   def check_opt(opts):
     prg = get_prg(opts=opts)

test/backend/test_multitensor.py

@@ -1,13 +1,13 @@
-import unittest, functools, random
+import unittest, random
 from tinygrad import Tensor, Device, nn, GlobalCounters, TinyJit, dtypes, Variable
 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 BufferCopy, CompiledRunner, run_schedule
+from tinygrad.engine.realize import CompiledRunner, run_schedule
 import numpy as np
 from hypothesis import given, strategies as strat, settings
-from test.helpers import not_support_multi_device, needs_second_gpu, slow
+from test.helpers import not_support_multi_device, needs_second_gpu, slow, call_is_graph
 
 settings.register_profile("my_profile", max_examples=200, deadline=None, derandomize=getenv("DERANDOMIZE_CI", False))
 settings.load_profile("my_profile")
@@ -544,7 +544,7 @@ class TestMultiTensor(unittest.TestCase):
       b.shard_(devices_2)
       c = jf(a, b)
       np.testing.assert_allclose(c.numpy(), a.numpy()+b.numpy(), atol=1e-4, rtol=1e-5)
-    assert len(jf.jit_cache) > 0
+    assert jf.captured is not None
 
   def test_multi_tensor_jit_body(self):
     @TinyJit
@@ -558,7 +558,7 @@ class TestMultiTensor(unittest.TestCase):
     for _ in range(5):
       r = jf()
       np.testing.assert_allclose(r.numpy(), np.ones(256)+np.ones(256), atol=1e-4, rtol=1e-5)
-    assert len(jf.jit_cache) > 0
+    assert jf.captured is not None
 
   def test_multitensor_jit_in_list(self):
     # test MULTI tensor inside a list container - exercises the container unpacking + MULTI unpacking
@@ -618,15 +618,12 @@ class TestMultiTensor(unittest.TestCase):
       o = jf(a, b, c, d).numpy()
       np.testing.assert_allclose(ref, o, atol=1e-4, rtol=1e-5)
 
-    graph_d0 = Device[d0].graph.func if isinstance(Device[d0].graph, functools.partial) else Device[d0].graph
-    graph_d1 = Device[d1].graph.func if isinstance(Device[d1].graph, functools.partial) else Device[d1].graph
     # Checking that 2 graphs per device, 1 copy and 1 last graph on device 1 are created.
-    assert isinstance(jf.jit_cache[0].prg, graph_d0)
-    assert isinstance(jf.jit_cache[1].prg, graph_d0)
-    assert isinstance(jf.jit_cache[2].prg, graph_d1)
-    assert isinstance(jf.jit_cache[3].prg, graph_d1)
-    assert isinstance(jf.jit_cache[4].prg, BufferCopy)
-    assert isinstance(jf.jit_cache[5].prg, graph_d1)
+    sis = jf.captured.linear.src
+    assert len(sis) == 6
+    for si in (sis[0], sis[1], sis[2], sis[3], sis[5]):
+      assert call_is_graph(si)
+    assert sis[4].src[0].op is Ops.COPY
 
   def test_bn_ast_on_devices(self):
     t = Tensor.empty((16, 64, 112, 112)).shard(devices_4, axis=0)

test/backend/test_ops.py

@@ -94,7 +94,7 @@ def prepare_test_op(low, high, shps, vals, forward_only=False):
 class TestOps(unittest.TestCase):
 
   def helper_test_exception(self, shps, torch_fxn, tinygrad_fxn=None, expected=None, forward_only=False, exact=False, vals=None, low=-1.5, high=1.5):
-    if getenv("MOCKGPU") and Device.DEFAULT == "NV": self.skipTest('helper_test_exception fails in CI CUDA')
+    if DEV.interface.startswith("MOCK") and Device.DEFAULT == "NV": self.skipTest('helper_test_exception fails in CI CUDA')
     ts, tst = prepare_test_op(low, high, shps, vals, forward_only)
     if tinygrad_fxn is None:
       tinygrad_fxn = torch_fxn
@@ -281,6 +281,17 @@ class TestOps(unittest.TestCase):
     helper_test_op([], lambda: torch.arange(5.5, 175.5, 2.5), lambda: Tensor.arange(5.5, 175.5, 2.5), forward_only=True)
     helper_test_op([], lambda: torch.arange(-30.2, -0.3, 0.75), lambda: Tensor.arange(-30.2, -0.3, 0.75), forward_only=True)
     helper_test_op([], lambda: torch.arange(-50.3, -380.2, -2.25), lambda: Tensor.arange(-50.3, -380.2, -2.25), forward_only=True)
+    # boundary values that fit exactly in int8 (min=-128, max=127)
+    helper_test_op([], lambda: torch.arange(128, dtype=torch.int8), lambda: Tensor.arange(128, dtype=dtypes.int8), forward_only=True)
+    helper_test_op([], lambda: torch.arange(-128, 128, dtype=torch.int8), lambda: Tensor.arange(-128, 128, dtype=dtypes.int8), forward_only=True)
+    helper_test_op([], lambda: torch.arange(127, -129, -1, dtype=torch.int8),
+                   lambda: Tensor.arange(127, -129, -1, dtype=dtypes.int8), forward_only=True)
+    # overflow: tinygrad raises (torch silently wraps)
+    with self.assertRaises(OverflowError): Tensor.arange(2**33, dtype=dtypes.int)
+    with self.assertRaises(OverflowError): Tensor.arange(129, dtype=dtypes.int8)            # last=128 overflows
+    with self.assertRaises(OverflowError): Tensor.arange(-129, 128, dtype=dtypes.int8)      # start=-129 overflows
+    with self.assertRaises(OverflowError): Tensor.arange(128, 0, -1, dtype=dtypes.int8)     # start=128 overflows
+    with self.assertRaises(OverflowError): Tensor.arange(127, -130, -1, dtype=dtypes.int8)  # last=-129 overflows
 
   def test_arange_big(self):
     helper_test_op([], lambda: torch.arange(256, dtype=torch.int32), lambda: Tensor.arange(256), forward_only=True)
@@ -866,7 +877,7 @@ class TestOps(unittest.TestCase):
     helper_test_op([(45,65)], lambda x: x.sin())
     helper_test_op([()], lambda x: x.sin())
     # works on real CUDA but not CI
-    if not ((getenv("MOCKGPU") and Device.DEFAULT == "NV") or Device.DEFAULT == "WEBGPU"):
+    if not ((DEV.interface.startswith("MOCK") and Device.DEFAULT == "NV") or Device.DEFAULT == "WEBGPU"):
       helper_test_op(None, lambda x: x.sin(), vals=[[math.nan, math.inf, -math.inf, 0.0]])
       helper_test_op(None, lambda x: x.sin(), vals=[[1e1, 1e2, 1e3, 1e4, 1e5, 1e6, -1e1, -1e2, -1e3, -1e4, -1e5, -1e6]],
                     atol=3e-3, rtol=3e-3, grad_atol=3e-3, grad_rtol=3e-3)
@@ -875,7 +886,7 @@ class TestOps(unittest.TestCase):
   def test_cos(self):
     helper_test_op([(45,65)], lambda x: x.cos())
     helper_test_op([()], lambda x: x.cos())
-    if not ((getenv("MOCKGPU") and Device.DEFAULT == "NV") or Device.DEFAULT == "WEBGPU"):
+    if not ((DEV.interface.startswith("MOCK") and Device.DEFAULT == "NV") or Device.DEFAULT == "WEBGPU"):
       helper_test_op(None, lambda x: x.cos(), vals=[[math.nan, math.inf, -math.inf, 0.0]])
       helper_test_op(None, lambda x: x.cos(), vals=[[1e1, 1e2, 1e3, 1e4, 1e5, 1e6, -1e1, -1e2, -1e3, -1e4, -1e5, -1e6]],
                     atol=3e-3, rtol=3e-3, grad_atol=3e-3, grad_rtol=3e-3)
@@ -886,7 +897,7 @@ class TestOps(unittest.TestCase):
     helper_test_op([(45,65)], lambda x: x.tan(), low=-1.5, high=1.5)
     helper_test_op([(45,65)], lambda x: x.tan(), low=-5, high=5)
     helper_test_op([()], lambda x: x.tan())
-    if not ((getenv("MOCKGPU") and Device.DEFAULT == "NV") or Device.DEFAULT == "WEBGPU"):
+    if not ((DEV.interface.startswith("MOCK") and Device.DEFAULT == "NV") or Device.DEFAULT == "WEBGPU"):
       helper_test_op(None, lambda x: x.tan(), vals=[[math.nan, math.inf, -math.inf, 0.0]])
       helper_test_op(None, lambda x: x.tan(), vals=[[1e1, 1e2, 1e3, 1e4, 1e5, 1e6, -1e1, -1e2, -1e3, -1e4, -1e5, -1e6]],
                     atol=3e-3, rtol=3e-3, grad_atol=3e-3, grad_rtol=3e-3)
@@ -3299,7 +3310,7 @@ class TestOps(unittest.TestCase):
     helper_test_op([(32,10)], lambda x: x.masked_fill((x>0.1).detach(), -math.inf))
     helper_test_op([(32,10)], lambda x: x.masked_fill((x<0.1).detach(), -math.inf))
 
-  @unittest.skipIf((getenv("MOCKGPU") or Device.DEFAULT == "PYTHON"), "very slow on MOCKGPU because reduce does not fold")
+  @unittest.skipIf((DEV.interface.startswith("MOCK") or Device.DEFAULT == "PYTHON"), "very slow on MOCKGPU because reduce does not fold")
   @unittest.skipIf(Device.DEFAULT == "WEBGPU", "webgpu runtime issue")
   @unittest.skipIf(Device.DEFAULT == "QCOM", "QCOM fails with: Resource deadlock avoided")
   def test_masked_select(self):

test/backend/test_opt_gemm.py

@@ -5,6 +5,7 @@ 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 test.helpers import replace_opts
 
 class TestOptGemm(unittest.TestCase):
   @classmethod
@@ -19,7 +20,7 @@ class TestOptGemm(unittest.TestCase):
     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(si.ast, renderer=Device[Device.DEFAULT].renderer, opts=opts))
+    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)
     np.testing.assert_allclose(self.res, test, atol=1e-4)

test/backend/test_pickle.py

@@ -125,6 +125,13 @@ class TestPickle(unittest.TestCase):
     out = add_fxn(x, y)
     np.testing.assert_equal(out.numpy(), 102)
 
+  def test_pickle_jit_no_del(self):
+    @TinyJit
+    def fn(x): return x + 1.0
+    for _ in range(3): fn(Tensor.randn(4))
+    loaded = pickle.loads(pickle.dumps(fn))
+    self.assertEqual(loaded(Tensor([1.0,2.0,3.0,4.0])).tolist(), [2.0,3.0,4.0,5.0])
+
   def test_pickle_context_var(self):
     v = ContextVar("test_var", 0)
     with Context(test_var=1):

test/backend/test_profiler.py

@@ -1,11 +1,11 @@
 import unittest, struct, contextlib, statistics, gc
 from tinygrad import Device, Tensor, dtypes, TinyJit
-from tinygrad.helpers import CI, getenv, Context, ProfileRangeEvent, cpu_profile, cpu_events, ProfilePointEvent, dedup
+from tinygrad.helpers import CI, DEV, Context, ProfileRangeEvent, cpu_profile, cpu_events, ProfilePointEvent, dedup
 from tinygrad.device import Buffer, BufferSpec, Compiled, ProfileDeviceEvent, ProfileGraphEvent
 from tinygrad.runtime.support.hcq import HCQCompiled
 from tinygrad.engine.realize import get_runner
 
-MOCKGPU = getenv("MOCKGPU")
+MOCKGPU = DEV.interface.startswith("MOCK")
 def _dev_base(d):
   p = d.split(":")
   return p[0] if len(p) < 2 or not p[1].isdigit() else f"{p[0]}:{p[1]}"

test/backend/test_quantize_onnx.py

@@ -7,6 +7,7 @@ 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 test.helpers import replace_opts
 
 N = 512
 
@@ -39,7 +40,7 @@ def create_gemm_model(model_path:str, batch_size=N, in_size=N, out_size=N, bias=
 
 def sexec(out:Tensor, opts:list[Opt], replace_src=None, run_count=3):
   si = out.schedule()[-1]
-  prg = get_program(si.ast, renderer=Device[Device.DEFAULT].renderer, opts=opts)
+  prg = get_program(replace_opts(si.ast, 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"))

test/backend/test_randomness.py

@@ -360,6 +360,7 @@ class TestRandomness(unittest.TestCase):
     torch_samples = [torch.tensor(w).multinomial(1, replacement=False).item() for _ in range(1000)]
     self.assertTrue(equal_distribution(lambda *_: Tensor(tiny_samples), lambda _: torch.tensor(torch_samples)))
 
+  @unittest.skip("this test is flaky")
   def test_multinomial_counterexample(self):
     tiny_res = Tensor([0.3, 0.6, 0.1]).multinomial(4000, replacement=True)
     torch_res = torch.tensor([0.3, 0.6, 0.1]).multinomial(4000, replacement=True)

test/backend/test_schedule.py

@@ -12,18 +12,20 @@ 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
+from tinygrad.engine.realize import CompiledRunner, run_schedule, run_linear
+from tinygrad.schedule import linear_to_schedule
 
 class KernelCountException(Exception): pass
 def check_schedule(t:Tensor|list[Tensor]|UOp, allowed:int, to_prerealize:list[Tensor]|None=None, filter_sink=True):
   if to_prerealize:
     with Context(DEBUG=0, TRACK_MATCH_STATS=0): Tensor.realize(*to_prerealize)
-  if isinstance(t, Tensor): sched = t.schedule()
-  elif isinstance(t, list) and isinstance(t[0], Tensor): sched = Tensor.schedule(*t)
+  if isinstance(t, Tensor): linear, var_vals = t.linear_with_vars()
+  elif isinstance(t, list) and isinstance(t[0], Tensor): linear, var_vals = Tensor.linear_with_vars(*t)
   else:
     assert isinstance(t, UOp), f"can't schedule {t}"
-    sched = Tensor(t).schedule()
+    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])
   if kernel_cnt != allowed:
@@ -33,7 +35,7 @@ def check_schedule(t:Tensor|list[Tensor]|UOp, allowed:int, to_prerealize:list[Te
         print("kernel", i+1)
         print(s.ast)
     raise KernelCountException(f"{kernel_cnt} != {allowed}")
-  return sched
+  return linear, var_vals
 
 def _realize_weights(m):
   for p in nn.state.get_parameters(m): p.realize()
@@ -91,17 +93,17 @@ class TestSchedule(unittest.TestCase):
   # all permutes, reshapes, expands and shrinks push through the reduce
   def test_arange_sum(self):
     a = Tensor.arange(6).reshape(3, 2).sum(axis=1)
-    run_schedule(check_schedule(a, 1))
+    run_linear(*check_schedule(a, 1))
     self.assertListEqual(a.tolist(), [1, 5, 9])
 
   def test_arange_sum_alt(self):
     a = (Tensor.arange(5).reshape(1,5).expand(6,5)*Tensor(2)).reshape(1,6,5).sum(axis=2)
-    run_schedule(check_schedule(a, 1))
+    run_linear(*check_schedule(a, 1))
     np.testing.assert_equal(a.numpy(), 20)
 
   def test_permute_arange(self):
     a = Tensor.arange(6).reshape(6, 1, 1).permute(2, 0, 1).sum(axis=1)
-    run_schedule(check_schedule(a, 1))
+    run_linear(*check_schedule(a, 1))
     self.assertListEqual(a.tolist(), [[15]])
 
   @unittest.skipUnless(is_dtype_supported(dtypes.half), "need half")
@@ -109,13 +111,13 @@ class TestSchedule(unittest.TestCase):
   def test_expand_buffer_before_cast(self):
     a = Tensor.randn(4, 2, 1).realize().permute((1, 0, 2))
     b = a.cast(dtypes.half).expand((2, 4, 4))+2
-    run_schedule(check_schedule(b, 1))
+    run_linear(*check_schedule(b, 1))
     np.testing.assert_allclose(b.numpy(), np.broadcast_to(a.numpy().astype(np.float16), (2, 4, 4))+2, rtol=1e-3)
 
   def test_indexing_scalars_simple(self):
     X = Tensor.randn(2, 2).realize()
     xt = X[Tensor(1)][Tensor(0)]
-    run_schedule(check_schedule(xt, 1))
+    run_linear(*check_schedule(xt, 1))
     np.testing.assert_equal(xt.numpy(), X.numpy()[1][0])
 
   @unittest.skipIf(CI and Device.DEFAULT == "NV", "crashes on NV CI")
@@ -135,46 +137,46 @@ class TestSchedule(unittest.TestCase):
     assume(a<x and b<y)
     X = Tensor.randn(x, y).realize()
     xt = X[Tensor(a)][Tensor(b)]
-    run_schedule(check_schedule(xt, 1))
+    run_linear(*check_schedule(xt, 1))
     np.testing.assert_equal(xt.numpy(), X.numpy()[a][b])
 
   def test_push_pads_elementwise(self):
     x = Tensor.full((4,4), 2.).contiguous().realize()
     y = Tensor.full((4,4), 4.).contiguous().realize()
     z = (x.reciprocal()*y).pad((None, (0,1),)).sum()
-    run_schedule(check_schedule(z, 1))
+    run_linear(*check_schedule(z, 1))
     self.assertEqual(z.item(), 32)
 
   def test_push_pads_contiguous(self):
     x = Tensor.full((4,1), 2.).contiguous()
     y = Tensor.full((4,4), 4.).contiguous()
     z = (x.reciprocal().expand(4,4)*y).pad((None, (0,1),)).sum()
-    run_schedule(check_schedule(z, 1, [x,y]))
+    run_linear(*check_schedule(z, 1, [x,y]))
     self.assertEqual(z.item(), 32)
 
   def test_constants_can_store(self):
     a = Tensor(2).contiguous()
-    run_schedule(check_schedule(a, 1))
+    run_linear(*check_schedule(a, 1))
     np.testing.assert_equal(a.numpy(), 2)
 
   def test_allow_push_permutes(self):
     a = Tensor.randn(10,10,10).realize()
     b = Tensor.randn(10,10,1).realize()
     c = a.sum(axis=0, keepdim=True).permute(2,1,0) + b
-    run_schedule(check_schedule(c, 1))
+    run_linear(*check_schedule(c, 1))
     np.testing.assert_allclose(c.numpy(), np.sum(a.numpy(), axis=0, keepdims=True).transpose(2,1,0)+b.numpy())
 
   def test_div_collapse_buffer(self):
     a = Tensor.full((4,), 4.0).contiguous().realize()
     b = Tensor.full((4,), 2.0).contiguous().realize()
     expr = (a*b)/b
-    run_schedule(check_schedule(expr, 1))
+    run_linear(*check_schedule(expr, 1))
     np.testing.assert_allclose(expr.numpy(), np.full((4,), 4.0))
 
   def test_div_collapse_const(self):
     a = Tensor.full((4,), 4.0).contiguous().realize()
     expr = a/a
-    run_schedule(check_schedule(expr, 1))
+    run_linear(*check_schedule(expr, 1))
     np.testing.assert_allclose(expr.numpy(), np.full((4,), 1.0))
 
   def test_div_collapse(self):
@@ -195,7 +197,7 @@ class TestSchedule(unittest.TestCase):
     a = Tensor.ones(4).contiguous()
     b = Tensor.ones(4).contiguous()
     sched = check_schedule([a, b], 1)
-    run_schedule(sched)
+    run_linear(*sched)
     # a and b share the same underlying device memory
     self.assertIs(a.uop.realized, b.uop.realized)
 
@@ -204,7 +206,7 @@ class TestSchedule(unittest.TestCase):
     a = src.clone()
     b = src.clone()
     sched = check_schedule([a, b], 2, filter_sink=False)
-    run_schedule(sched)
+    run_linear(*sched)
     # a and b are assigned to the same device Buffer
     self.assertIsNot(a.uop.base.realized, b.uop.base.realized)
 
@@ -226,7 +228,7 @@ class TestSchedule(unittest.TestCase):
     ax = r.reshape(1)*2
     ay = r.reshape(1).shrink(((1,1),))*2
     out = ax+ay.pad(((1, 0),))
-    run_schedule(check_schedule(out, 1))
+    run_linear(*check_schedule(out, 1))
     self.assertEqual(out.item(), 4.)
 
   def test_preserve_multistage_reduce(self):
@@ -234,7 +236,7 @@ class TestSchedule(unittest.TestCase):
     x = Tensor.randn(big_enough).realize()
     with Context(SPLIT_REDUCEOP=1):
       out = (x - x.max(keepdim=True)).max()
-      run_schedule(check_schedule(out, 4))
+      run_linear(*check_schedule(out, 4))
     np.testing.assert_allclose(out.numpy(), (x.numpy() - x.numpy().max(keepdims=True)).max())
 
   @unittest.skip("these two Tensors are the same")
@@ -244,7 +246,7 @@ class TestSchedule(unittest.TestCase):
     z = y.matmul(x).sum()
     z.backward()
     out = x.grad.contiguous()
-    run_schedule(check_schedule(out, 1))
+    run_linear(*check_schedule(out, 1))
     np.testing.assert_allclose(out.numpy(), np.ones((64,64)))
 
   def test_example_matmul_contig(self):
@@ -253,7 +255,7 @@ class TestSchedule(unittest.TestCase):
     z = y.matmul(x).sum()
     z.backward()
     out = x.grad.contiguous()
-    run_schedule(check_schedule(out, 1))
+    run_linear(*check_schedule(out, 1))
     np.testing.assert_allclose(out.numpy(), np.ones((64,64)))
 
   def test_example_matmul_same(self):
@@ -261,7 +263,7 @@ class TestSchedule(unittest.TestCase):
     z = x.matmul(x).sum()
     z.backward()
     out = x.grad.contiguous()
-    run_schedule(check_schedule(out, 1))
+    run_linear(*check_schedule(out, 1))
     # NOTE: the gradient flows twice
     np.testing.assert_allclose(out.numpy(), 2*np.ones((64,64)))
 
@@ -275,7 +277,7 @@ class TestSchedule(unittest.TestCase):
     b_out = b.sum(1)
     b_out = b_out[:16]
     out = a_out + b_out + c
-    run_schedule(check_schedule(out, 1))
+    run_linear(*check_schedule(out, 1))
     np.testing.assert_allclose(out.numpy(), a.numpy().sum(axis=1)[:16] + b.numpy().sum(axis=1)[:16] + c.numpy(), atol=1e-4, rtol=1e-4)
 
   def test_reduce_same_size(self):
@@ -284,7 +286,7 @@ class TestSchedule(unittest.TestCase):
     out0 = a.sum() + 2
     out1 = a.sum() + 4
     out2 = out0 * out1
-    run_schedule(check_schedule([out0, out1, out2], 3)) # TODO: 1?
+    run_linear(*check_schedule([out0, out1, out2], 3)) # TODO: 1?
     np.testing.assert_allclose(out0.numpy(), out0_np:=a.numpy().sum()+2, atol=1e-4, rtol=1e-6)
     np.testing.assert_allclose(out1.numpy(), out1_np:=a.numpy().sum()+4, atol=1e-4, rtol=1e-6)
     np.testing.assert_allclose(out2.numpy(), out0_np*out1_np, atol=1e-4, rtol=1e-6)
@@ -295,7 +297,7 @@ class TestSchedule(unittest.TestCase):
     out0 = a.sum().exp2()
     # out1 has two paths to a.sum()
     out1 = a.sum() + out0
-    run_schedule(check_schedule([out0, out1], 2)) # TODO: 1?
+    run_linear(*check_schedule([out0, out1], 2)) # TODO: 1?
     np.testing.assert_allclose(out0.numpy(), out0_np:=np.exp2(a.numpy().sum()), atol=1e-4, rtol=1e-4)
     np.testing.assert_allclose(out1.numpy(), a.numpy().sum()+out0_np, atol=1e-4, rtol=1e-6)
 
@@ -307,8 +309,8 @@ class TestSchedule(unittest.TestCase):
     b = (a + out0 + out1)
     out2 = b.sum().exp2()
     out3 = b.sum() + out2
-    # run_schedule(check_schedule([out0, out1, out2, out3], 1))
-    run_schedule(check_schedule([out0, out1, out2, out3], 4))
+    # run_linear(*check_schedule([out0, out1, out2, out3], 1))
+    run_linear(*check_schedule([out0, out1, out2, out3], 4))
     np.testing.assert_allclose(out0.numpy(), np_out0:=np.exp2(a.numpy().sum()), atol=1e-4, rtol=1e-4)
     np.testing.assert_allclose(out1.numpy(), np_out1:=a.numpy().sum()+np_out0, atol=1e-4, rtol=1e-4)
     np_b = (a.numpy() + np_out0 + np_out1)
@@ -323,8 +325,8 @@ class TestSchedule(unittest.TestCase):
     # b.sum() is not a descendant of the fused nodes
     out0 = a.sum() + b.sum() + 2
     out1 = a.sum() + b.sum() + 4
-    # run_schedule(check_schedule([out0, out1], 1))
-    run_schedule(check_schedule([out0, out1], 2))
+    # run_linear(*check_schedule([out0, out1], 1))
+    run_linear(*check_schedule([out0, out1], 2))
     np.testing.assert_allclose(out0.numpy(), a.numpy().sum()+b.numpy().sum()+2, atol=1e-4, rtol=1e-4)
     np.testing.assert_allclose(out1.numpy(), a.numpy().sum()+b.numpy().sum()+4, atol=1e-4, rtol=1e-4)
 
@@ -336,8 +338,8 @@ class TestSchedule(unittest.TestCase):
     # reduce node in the indirect path from r to out2
     out1 = (a - out0).max()
     out2 = r + out1
-    # run_schedule(check_schedule([r, out0, out1, out2], 1))
-    run_schedule(check_schedule([r, out0, out1, out2], 4))
+    # run_linear(*check_schedule([r, out0, out1, out2], 1))
+    run_linear(*check_schedule([r, out0, out1, out2], 4))
     np.testing.assert_allclose(r.numpy(), r_np:=a.numpy().sum(), atol=1e-4, rtol=1e-4)
     np.testing.assert_allclose(out0.numpy(), out0_np:=np.exp2(r_np), atol=1e-4, rtol=1e-4)
     np.testing.assert_allclose(out1.numpy(), out1_np:=(a.numpy() - out0_np).max(), atol=1e-4, rtol=1e-4)
@@ -350,8 +352,8 @@ class TestSchedule(unittest.TestCase):
     out0 = a.sum() + 4
     out1 = b.max() + out0*2
     out2 = a.sum() + out1
-    # run_schedule(check_schedule([out0, out1, out2], 1))
-    run_schedule(check_schedule([out0, out1, out2], 3))
+    # run_linear(*check_schedule([out0, out1, out2], 1))
+    run_linear(*check_schedule([out0, out1, out2], 3))
     np.testing.assert_allclose(out0.numpy(), out0_np:=a.numpy().sum()+4, atol=1e-4, rtol=1e-6)
     np.testing.assert_allclose(out1.numpy(), out1_np:=b.numpy().max() + out0_np*2, atol=1e-4, rtol=1e-6)
     np.testing.assert_allclose(out2.numpy(), a.numpy().sum() + out1_np, atol=1e-4, rtol=1e-6)
@@ -365,8 +367,8 @@ class TestSchedule(unittest.TestCase):
     # e1 is in the indirect path from a.sum() to out1
     e = b + out0
     out1 = r + e[0][0][0]
-    # run_schedule(check_schedule([r, out0, out1, e], 3)) # 1 or 2 or 3? should be 1 (one reduce) but the different outputs might make it 3
-    run_schedule(check_schedule([r, out0, out1, e], 4))
+    # run_linear(*check_schedule([r, out0, out1, e], 3)) # 1 or 2 or 3? should be 1 (one reduce) but the different outputs might make it 3
+    run_linear(*check_schedule([r, out0, out1, e], 4))
     np.testing.assert_allclose(r.numpy(), r_np:=a.numpy().sum(), atol=1e-4, rtol=1e-4)
     np.testing.assert_allclose(out0.numpy(), out0_np:=np.exp2(r_np), atol=1e-4, rtol=1e-4)
     np.testing.assert_allclose(e.numpy(), e_np:=b.numpy() + out0_np, atol=1e-4, rtol=1e-4)
@@ -378,8 +380,8 @@ class TestSchedule(unittest.TestCase):
     b = Tensor.randn((1, 16)).realize()
     out0 = a.sum() + 2
     out1 = a.sum() + b
-    # run_schedule(check_schedule([out0, out1], 2))
-    run_schedule(check_schedule([out0, out1], 3))
+    # run_linear(*check_schedule([out0, out1], 2))
+    run_linear(*check_schedule([out0, out1], 3))
     np.testing.assert_allclose(out0.numpy(), a.numpy().sum()+2, atol=1e-4, rtol=1e-4)
     np.testing.assert_allclose(out1.numpy(), a.numpy().sum()+b.numpy(), atol=1e-4, rtol=1e-4)
 
@@ -387,7 +389,7 @@ class TestSchedule(unittest.TestCase):
     Tensor.manual_seed(0)
     x = Tensor.randn(4, 32).realize()
     out = x.std(-1)
-    run_schedule(check_schedule(out, 2))
+    run_linear(*check_schedule(out, 2))
     np.testing.assert_allclose(out.numpy(), x.numpy().std(axis=-1, ddof=1), atol=1e-4, rtol=1e-4)
 
   def test_scaled_dot_product_attention_multireduce_fusion(self):
@@ -396,14 +398,14 @@ class TestSchedule(unittest.TestCase):
     k = Tensor.randn(32,8,16,8).realize()
     v = Tensor.randn(32,8,16,8).realize()
     out = Tensor.scaled_dot_product_attention(q,k,v)
-    run_schedule(check_schedule(out, 4))
+    run_linear(*check_schedule(out, 4))
     if getenv("CHECK", 1):
       import torch
       compare = torch.nn.functional.scaled_dot_product_attention(torch.tensor(q.numpy()),torch.tensor(k.numpy()),torch.tensor(v.numpy()))
       np.testing.assert_allclose(out.numpy(), compare.numpy(), atol=1e-6, rtol=1e-3)
 
     out = Tensor.scaled_dot_product_attention(q,k,v)
-    run_schedule(check_schedule(out, 4)) # TODO: should be 1?
+    run_linear(*check_schedule(out, 4)) # TODO: should be 1?
     if getenv("CHECK", 1):
       import torch
       compare = torch.nn.functional.scaled_dot_product_attention(torch.tensor(q.numpy()),torch.tensor(k.numpy()),torch.tensor(v.numpy()))
@@ -415,8 +417,8 @@ class TestSchedule(unittest.TestCase):
     b = Tensor.randn(4, 32).realize()
     c = Tensor.randn(4, 32).realize()
     out = (c * a.sum(-1, keepdim=True)).sum(-1) + (b * a.sum(-1, keepdim=True)).sum(-1) # a.sum has >1 children but should still fuse
-    # run_schedule(check_schedule(out, 1))
-    run_schedule(check_schedule(out, 2))
+    # run_linear(*check_schedule(out, 1))
+    run_linear(*check_schedule(out, 2))
     np.testing.assert_allclose(out.numpy(), \
       (c.numpy()*a.numpy().sum(axis=-1,keepdims=True)).sum(-1) + (b.numpy()*a.numpy().sum(axis=-1,keepdims=True)).sum(-1), atol=1e-4, rtol=1e-4)
 
@@ -424,16 +426,16 @@ class TestSchedule(unittest.TestCase):
     Tensor.manual_seed(0)
     a = Tensor.randn(4, 32).realize()
     out = (a+a.sum(-1, keepdim=True)).sum(-1)
-    # run_schedule(check_schedule(out, 1))
-    run_schedule(check_schedule(out, 2))
+    # run_linear(*check_schedule(out, 1))
+    run_linear(*check_schedule(out, 2))
     np.testing.assert_allclose(out.numpy(), (a.numpy()+a.numpy().sum(axis=-1,keepdims=True)).sum(axis=-1), atol=1e-4, rtol=1e-4)
 
   def test_reduce_expand_reduce_expand_fusion(self):
     Tensor.manual_seed(0)
     a = Tensor.randn(4, 32).realize()
     out = a+(a+a.sum(-1,keepdim=True)).sum(-1, keepdim=True)
-    # run_schedule(check_schedule(out, 2))
-    run_schedule(check_schedule(out, 3))
+    # run_linear(*check_schedule(out, 2))
+    run_linear(*check_schedule(out, 3))
     np.testing.assert_allclose(out.numpy(), \
       a.numpy()+(a.numpy()+a.numpy().sum(axis=-1,keepdims=True)).sum(axis=-1,keepdims=True), atol=1e-4, rtol=1e-4)
 
@@ -442,8 +444,8 @@ class TestSchedule(unittest.TestCase):
     a = Tensor.randn(4, 32).realize()
     out0 = a+a.sum(-1, keepdim=True)
     out1 = out0.sum(-1)
-    # run_schedule(check_schedule(out, 2))
-    run_schedule(check_schedule([out0, out1], 3))
+    # run_linear(*check_schedule(out, 2))
+    run_linear(*check_schedule([out0, out1], 3))
     np.testing.assert_allclose(out0.numpy(), a.numpy()+a.numpy().sum(axis=-1,keepdims=True), atol=1e-4, rtol=1e-4)
     np.testing.assert_allclose(out1.numpy(), (a.numpy()+a.numpy().sum(axis=-1,keepdims=True)).sum(axis=-1), atol=1e-4, rtol=1e-4)
 
@@ -452,8 +454,8 @@ class TestSchedule(unittest.TestCase):
     x = Tensor.randn(4, 32).realize()
     y = Tensor.randn(4, 32).realize()
     out = (y + x.sum(axis=-1, keepdim=True)).sum(axis=-1)
-    # run_schedule(check_schedule(out, 1))
-    run_schedule(check_schedule(out, 2))
+    # run_linear(*check_schedule(out, 1))
+    run_linear(*check_schedule(out, 2))
     np.testing.assert_allclose(out.numpy(), (y.numpy() + x.numpy().sum(axis=-1, keepdims=True)).sum(axis=-1), atol=1e-4, rtol=1e-4)
 
   def test_multireduce_fusion_simple_parallel(self):
@@ -461,15 +463,15 @@ class TestSchedule(unittest.TestCase):
     x = Tensor.randn(4, 32).realize()
     y = Tensor.randn(4, 32).realize()
     out = y.sum(axis=-1) + x.sum(axis=-1)
-    run_schedule(check_schedule(out, 1))
+    run_linear(*check_schedule(out, 1))
     np.testing.assert_allclose(out.numpy(), y.numpy().sum(axis=-1) + x.numpy().sum(axis=-1), atol=1e-4, rtol=1e-4)
 
   def test_multireduce_fusion_sequential(self):
     Tensor.manual_seed(0)
     x = Tensor.randn(4, 32).realize()
     out = x.std(-1)
-    # run_schedule(check_schedule(out, 1))
-    run_schedule(check_schedule(out, 2))
+    # run_linear(*check_schedule(out, 1))
+    run_linear(*check_schedule(out, 2))
     np.testing.assert_allclose(out.numpy(), x.numpy().std(axis=-1, ddof=1), atol=1e-4, rtol=1e-4)
 
   def test_multireduce_fusion_parallel(self):
@@ -477,16 +479,16 @@ class TestSchedule(unittest.TestCase):
     x = Tensor.randn(4, 32).realize()
     y = Tensor.randn(4, 32).realize()
     out = x.std(-1) + y.std(-1)
-    # run_schedule(check_schedule(out, 1))
-    run_schedule(check_schedule(out, 3))
+    # run_linear(*check_schedule(out, 1))
+    run_linear(*check_schedule(out, 3))
     np.testing.assert_allclose(out.numpy(), x.numpy().std(axis=-1, ddof=1) + y.numpy().std(axis=-1, ddof=1), atol=1e-4, rtol=1e-4)
 
   def test_multireduce_diffops_sequential(self):
     Tensor.manual_seed(0)
     x = Tensor.randn(4, 32).realize()
     out = (x - x.max(-1, keepdim=True)).sum(-1)
-    # run_schedule(check_schedule(out, 1))
-    run_schedule(check_schedule(out, 2))
+    # run_linear(*check_schedule(out, 1))
+    run_linear(*check_schedule(out, 2))
     np.testing.assert_allclose(out.numpy(), (x.numpy() - x.numpy().max(axis=-1, keepdims=True)).sum(axis=-1), atol=1e-4, rtol=1e-4)
 
   def test_multireduce_fusion_diffops_parallel(self):
@@ -494,7 +496,7 @@ class TestSchedule(unittest.TestCase):
     x = Tensor.randn(4, 32).realize()
     y = Tensor.randn(4, 32).realize()
     out = x.sum(-1) + y.max(-1)
-    run_schedule(check_schedule(out, 1))
+    run_linear(*check_schedule(out, 1))
     np.testing.assert_allclose(out.numpy(), x.numpy().sum(axis=-1) + y.numpy().max(axis=-1), atol=1e-4, rtol=1e-4)
 
   def test_multireduce_fusion_sequential_and_parallel(self):
@@ -505,8 +507,8 @@ class TestSchedule(unittest.TestCase):
     out = [((x - mu).square().sum(-1)/x.shape[-1]).sqrt(), ((y - mu).square().sum(-1)/y.shape[-1]).sqrt()]
     np_mu = (x.numpy() - x.numpy().max(axis=-1, keepdims=True)).mean(axis=-1, keepdims=True) + \
       (y.numpy() - y.numpy().max(axis=-1, keepdims=True)).mean(axis=-1, keepdims=True)
-    # run_schedule(check_schedule(out, 1))
-    run_schedule(check_schedule(out, 5))
+    # run_linear(*check_schedule(out, 1))
+    run_linear(*check_schedule(out, 5))
     np.testing.assert_allclose(out[0].numpy(), np.sqrt(np.square(x.numpy() - np_mu).sum(-1)/x.shape[-1]), atol=1e-4, rtol=1e-4)
     np.testing.assert_allclose(out[1].numpy(), np.sqrt(np.square(y.numpy() - np_mu).sum(-1)/y.shape[-1]), atol=1e-4, rtol=1e-4)
 
@@ -533,14 +535,14 @@ class TestSchedule(unittest.TestCase):
     a,b = Tensor.randn(4, 64).realize(), Tensor.rand(64,8).realize()
     c,d = Tensor.randn(4, 64).realize(), Tensor.rand(64,8).realize()
     out = a@b + c@d
-    run_schedule(check_schedule(out, 1))
+    run_linear(*check_schedule(out, 1))
     np.testing.assert_allclose(out.numpy(), a.numpy()@b.numpy() + c.numpy()@d.numpy(), atol=1e-4, rtol=1e-4)
 
   def test_softmax_fusion(self):
     Tensor.manual_seed(0)
     x = Tensor.randn(4, 12, 64, 64).realize()
     out = x.softmax()
-    run_schedule(check_schedule(out, 3))
+    run_linear(*check_schedule(out, 3))
     expected = (x_exp:=np.exp(x.numpy()-x.numpy().max(-1, keepdims=True)))/x_exp.sum(-1, keepdims=True)
     np.testing.assert_allclose(out.numpy(), expected, atol=1e-4, rtol=1e-4)
 
@@ -551,8 +553,8 @@ class TestSchedule(unittest.TestCase):
     layer.bias = Tensor.randn(10,10).realize()
     x = Tensor.randn(20, 5, 10, 10).realize()
     out = layer(x)
-    # run_schedule(check_schedule(out, 2))
-    run_schedule(check_schedule(out, 3))
+    # run_linear(*check_schedule(out, 2))
+    run_linear(*check_schedule(out, 3))
     y = (x.numpy() - x.numpy().mean(layer.axis, keepdims=True))
     expected = y / np.sqrt((y*y).mean(layer.axis, keepdims=True) + layer.eps)
     np.testing.assert_allclose(out.numpy(), expected * layer.weight.numpy() + layer.bias.numpy(), atol=1e-4, rtol=1e-4)
@@ -567,7 +569,7 @@ class TestSchedule(unittest.TestCase):
     # schedule = check_schedule([b,c], 3)
     # self.assertIs(schedule[0].ast[0].src[0].arg, Ops.MUL)
     schedule = check_schedule([b,c], 4)
-    run_schedule(schedule)
+    run_linear(*schedule)
     np.testing.assert_allclose(b.numpy(), np_r.sum(0) + 8, atol=1e-4, rtol=1e-4)
     np.testing.assert_allclose(c.numpy(), np_r.sum(1) + 12, atol=1e-4, rtol=1e-4)
 
@@ -579,7 +581,7 @@ class TestSchedule(unittest.TestCase):
     d = r.T * 4
     e = r * (d + a).sum(2)
     schedule = check_schedule([d, e], 3) # make sure it doesn't fuse
-    run_schedule(schedule)
+    run_linear(*schedule)
     np.testing.assert_allclose(d.numpy(), (a.numpy().sum(2) + b.numpy()).T * 4, atol=1e-4, rtol=1e-4)
     np.testing.assert_allclose(e.numpy(), (a.numpy().sum(2) + b.numpy()) * (d.numpy() + a.numpy()).sum(2), atol=1e-4, rtol=1e-4)
 
@@ -592,7 +594,7 @@ class TestSchedule(unittest.TestCase):
     r = a.sum(1) + c
     out = r[:4] * b + d.sum(1)[:4]
     schedule = check_schedule(out, 1)
-    run_schedule(schedule)
+    run_linear(*schedule)
     np.testing.assert_allclose(out.numpy(), (a.numpy().sum(1) + c.numpy())[:4] * b.numpy() + d.numpy().sum(1)[:4], atol=1e-4, rtol=1e-4)
 
   def test_multireduce_midreduce_nochase(self):
@@ -600,7 +602,7 @@ class TestSchedule(unittest.TestCase):
     a = Tensor.randn(16, 16).realize()
     b = (a.sum(0)+a.max(0) + a.max(1)+a.sum(1)) + 2
     schedule = check_schedule(b, 1)
-    run_schedule(schedule)
+    run_linear(*schedule)
     np.testing.assert_allclose(b.numpy(), a.numpy().sum(0)+a.numpy().max(0) + a.numpy().max(1)+a.numpy().sum(1)+2, atol=1e-4, rtol=1e-4)
 
   # pattern in test_transformer
@@ -610,8 +612,8 @@ class TestSchedule(unittest.TestCase):
     b = Tensor.randn(16, 16).realize()
     c = a.sum() + 2
     d = (a.sum() - b.sum()) * 4
-    # run_schedule(check_schedule([c, d], 1))
-    run_schedule(check_schedule([c, d], 2))
+    # run_linear(*check_schedule([c, d], 1))
+    run_linear(*check_schedule([c, d], 2))
     np.testing.assert_allclose(c.numpy(), a.numpy().sum()+2, atol=1e-4, rtol=1e-4)
     np.testing.assert_allclose(d.numpy(), (a.numpy().sum() - b.numpy().sum()) * 4, atol=1e-4, rtol=1e-4)
 
@@ -622,8 +624,8 @@ class TestSchedule(unittest.TestCase):
     b = Tensor.randn(16, 16).realize()
     c = a.sum() + 2
     d = b.sum() - c
-    # run_schedule(check_schedule([c, d], 1))
-    run_schedule(check_schedule([c, d], 2))
+    # run_linear(*check_schedule([c, d], 1))
+    run_linear(*check_schedule([c, d], 2))
     np.testing.assert_allclose(c.numpy(), a.numpy().sum()+2, atol=1e-4, rtol=1e-4)
     np.testing.assert_allclose(d.numpy(), b.numpy().sum()-(a.numpy().sum()+2), atol=1e-4, rtol=1e-4)
 
@@ -636,8 +638,8 @@ class TestSchedule(unittest.TestCase):
     d = a.sum() * 2
     e = c * d
     f = b.sum() - e
-    # run_schedule(check_schedule([c, d, e, f], 1))
-    run_schedule(check_schedule([c, d, e, f], 4))
+    # run_linear(*check_schedule([c, d, e, f], 1))
+    run_linear(*check_schedule([c, d, e, f], 4))
     np.testing.assert_allclose(c.numpy(), c_np:=a.numpy().sum()+2, atol=1e-4, rtol=1e-4)
     np.testing.assert_allclose(d.numpy(), d_np:=a.numpy().sum()*2, atol=1e-4, rtol=1e-4)
     np.testing.assert_allclose(e.numpy(), e_np:=c_np*d_np, atol=1e-4, rtol=1e-4)
@@ -651,8 +653,8 @@ class TestSchedule(unittest.TestCase):
     d = a.sum() * 2
     e = c * d
     f = (b - d).sum() - e
-    # run_schedule(check_schedule([c, d, e, f], 1))
-    run_schedule(check_schedule([c, d, e, f], 4))
+    # run_linear(*check_schedule([c, d, e, f], 1))
+    run_linear(*check_schedule([c, d, e, f], 4))
     np.testing.assert_allclose(c.numpy(), c_np:=a.numpy().sum()+2, atol=1e-4, rtol=1e-4)
     np.testing.assert_allclose(d.numpy(), d_np:=a.numpy().sum()*2, atol=1e-4, rtol=1e-4)
     np.testing.assert_allclose(e.numpy(), e_np:=c_np*d_np, atol=1e-4, rtol=1e-4)
@@ -663,7 +665,7 @@ class TestSchedule(unittest.TestCase):
     a = Tensor.rand(3, 4, 5).realize()
     b = Tensor.rand(3, 4, 5).realize()
     out = (a + b).pad(((0, 1), (0, 1), (0, 1)), value=1.0).sum().contiguous()
-    run_schedule(check_schedule(out, 1))
+    run_linear(*check_schedule(out, 1))
     np.testing.assert_allclose(out.numpy(), np.pad(a.numpy()+b.numpy(), ((0, 1), (0, 1), (0, 1)), constant_values=1.0).sum(), atol=1e-5, rtol=1e-6)
 
   def test_multireduce_pad_reduce_safe(self):
@@ -671,7 +673,7 @@ class TestSchedule(unittest.TestCase):
     a = Tensor.randn(3, 4, 5).realize()
     b = Tensor.randn(3, 4, 5).realize()
     out = (a.pad(((0, 1), (0, 1), (0, 1)), value=1.0).sum(keepdim=True)+b.pad(((0, 1), (0, 1), (0, 1)), value=1.0).sum()).contiguous()
-    run_schedule(check_schedule(out, 1))
+    run_linear(*check_schedule(out, 1))
     np.testing.assert_allclose(out.numpy(), np.pad(a.numpy(), ((0, 1), (0, 1), (0, 1)), constant_values=1.0).sum(keepdims=True) + \
                                                    np.pad(b.numpy(), ((0, 1), (0, 1), (0, 1)), constant_values=1.0).sum(), atol=1e-4, rtol=1e-4)
 
@@ -679,7 +681,7 @@ class TestSchedule(unittest.TestCase):
     Tensor.manual_seed(0)
     a = Tensor.rand(3, 4, 5).realize()
     out = a.log2().pad(((0, 1), (0, 1), (0, 1)), value=1.0).sum().contiguous()
-    run_schedule(check_schedule(out, 1))
+    run_linear(*check_schedule(out, 1))
     np.testing.assert_allclose(out.numpy(), np.pad(np.log2(a.numpy()), ((0, 1), (0, 1), (0, 1)), constant_values=1.0).sum(), atol=1e-5, rtol=1e-6)
 
   def test_multireduce_pad_reduce_unsafe(self):
@@ -687,8 +689,8 @@ class TestSchedule(unittest.TestCase):
     a = Tensor.randn(3, 4, 5).abs().realize()
     b = Tensor.randn(3, 4, 5).abs().realize()
     out = (a.log2().pad(((0, 1), (0, 1), (0, 1)), value=1.0).sum()+b).abs().log2().pad(((0, 1), (0, 1), (0, 1)), value=1.0).sum().contiguous()
-    # run_schedule(check_schedule(out, 1))
-    run_schedule(check_schedule(out, 2))
+    # run_linear(*check_schedule(out, 1))
+    run_linear(*check_schedule(out, 2))
     np.testing.assert_allclose(out.numpy(), np.pad(np.log2(np.abs(np.pad(np.log2(a.numpy()), ((0, 1), (0, 1), (0, 1)), constant_values=1.0).sum() + \
                                                    b.numpy())), ((0, 1), (0, 1), (0, 1)), constant_values=1.0).sum(), atol=3e-4, rtol=1e-5)
 
@@ -696,13 +698,13 @@ class TestSchedule(unittest.TestCase):
     a = Tensor.ones((3, )).contiguous().realize()
     b = Tensor.ones((3, )).contiguous().realize()
     out = (a + b).shrink(((0, 1),)).pad(((0, 1),)).contiguous()
-    run_schedule(check_schedule(out, 1))
+    run_linear(*check_schedule(out, 1))
     np.testing.assert_equal(out.numpy(), [2, 0])
 
   def test_shrink_pad_unsafe(self):
     a = Tensor.ones((3, )).contiguous().realize()
     out = a.exp2().shrink(((0, 1),)).pad(((0, 1),)).contiguous()
-    run_schedule(check_schedule(out, 1))
+    run_linear(*check_schedule(out, 1))
     np.testing.assert_equal(out.numpy(), [2, 0])
 
   def test_base_change_shrink_pad(self):
@@ -710,7 +712,7 @@ class TestSchedule(unittest.TestCase):
     b = a.exp2()
     c = b[:-1, :-1]
     d = c.pad(((0, 1), (0, 1))) * 2
-    run_schedule(check_schedule(d, 1))
+    run_linear(*check_schedule(d, 1))
     np.testing.assert_equal(d.numpy(), np.pad(np.exp2(a.numpy())[:-1, :-1], ((0, 1), (0, 1)))*2)
 
   def test_base_change_expand_pad(self):
@@ -718,7 +720,7 @@ class TestSchedule(unittest.TestCase):
     b = a.exp2()
     c = b[:, None, :]
     d = c.pad(((0, 0), (1, 1), (0, 0))) * 2
-    run_schedule(check_schedule(d, 1))
+    run_linear(*check_schedule(d, 1))
     np.testing.assert_equal(d.numpy(), np.pad(np.exp2(a.numpy())[:, None, :], ((0, 0), (1, 1), (0, 0)))*2)
 
   def test_fuse_arange_pad_replicate_mode(self):
@@ -726,7 +728,7 @@ class TestSchedule(unittest.TestCase):
     y = x.pad((-1,2,2,-1), mode="replicate")
     dx = y.sum().gradient(x)[0]
     sched = check_schedule(dx, 1)
-    run_schedule(sched)
+    run_linear(*sched)
     np.testing.assert_allclose(dx.numpy(), [[[[0.,3.,9.],[0,1.,3.],[0.,0.,0.]]]*3]*3)
 
   # TODO like openpilot with imagef
@@ -735,7 +737,7 @@ class TestSchedule(unittest.TestCase):
     a = Tensor.ones(4, 4).contiguous().realize()
     b = a.cast(dtypes.half).expand(2, 4, 4)
     c = b.cast(dtypes.int).expand(2, 2, 4, 4)
-    run_schedule(check_schedule(c, 1))
+    run_linear(*check_schedule(c, 1))
     np.testing.assert_equal(c.numpy(), np.ones(((2, 2, 4, 4)), dtype=np.int32))
 
   def test_base_change_pad_expand(self):
@@ -743,7 +745,7 @@ class TestSchedule(unittest.TestCase):
     b = Tensor.full((4, 4), 2.).contiguous().realize()
     c = (a + b).pad(((1, 1), (1, 1)))
     d = c.cast(dtypes.int).expand((2, 6, 6)) * 4
-    run_schedule(check_schedule(d, 1))
+    run_linear(*check_schedule(d, 1))
     c_np = np.pad((np.full((4, 4), 2., dtype=np.float32) + np.full((4, 4), 1., dtype=np.float32)), ((1, 1), (1, 1)), constant_values=0.0)
     np.testing.assert_equal(d.numpy(), np.broadcast_to(c_np.astype(np.half), (2, *c_np.shape)) * 4)
 
@@ -754,7 +756,7 @@ class TestSchedule(unittest.TestCase):
     p = P[0]
     p = p.pad(((1, 0), ))
     p = p.repeat([2])
-    run_schedule(check_schedule(p, 4))  # TODO: this is high
+    run_linear(*check_schedule(p, 3))
     tiny_ret = p.numpy()
 
     P = np.ones((3, 3), dtype=np.float32)
@@ -836,7 +838,7 @@ class TestSchedule(unittest.TestCase):
 
   def _test_fusion(self, shapes, f, cnt):
     with Context(DEBUG=0, TRACK_MATCH_STATS=0): args = [Tensor.randn(s).realize() for s in shapes]
-    run_schedule(check_schedule(compare:=f(*args), cnt))
+    run_linear(*check_schedule(compare:=f(*args), cnt))
     if getenv("COMPARE", 1):
       import torch
       good = f(*[torch.tensor(x.numpy()) for x in args])
@@ -880,9 +882,9 @@ class TestSchedule(unittest.TestCase):
   def test_cast_const_view(self):
     a = Tensor.ones((4, 4), dtype=dtypes.float32)
     casted_view = a.cast(dtypes.int32)
-    run_schedule(check_schedule(casted_view, 1))
+    run_linear(*check_schedule(casted_view, 1))
     realized_const_view = casted_view.contiguous()
-    run_schedule(check_schedule(realized_const_view, 0))
+    run_linear(*check_schedule(realized_const_view, 0))
     self.assertListEqual(realized_const_view.tolist(), [[1, 1, 1, 1], [1, 1, 1, 1], [1, 1, 1, 1], [1, 1, 1, 1]])
 
   @given(strat.sampled_from(dtypes.all), strat.sampled_from(dtypes.all))
@@ -891,41 +893,41 @@ class TestSchedule(unittest.TestCase):
     assume(is_dtype_supported(dt1) and is_dtype_supported(dt2))
     a = Tensor(1, dtype=dt1).reshape(1, 1).pad(((1, 1), None))
     casted_view = a.cast(dt2)
-    run_schedule(check_schedule(casted_view, 0))
+    run_linear(*check_schedule(casted_view, 0))
     realized_const_view = casted_view.contiguous()
-    run_schedule(check_schedule(realized_const_view, 1))
+    run_linear(*check_schedule(realized_const_view, 1))
     np.testing.assert_equal(realized_const_view.numpy(), [[0], [1], [0]])
 
   def test_simple_indexing(self):
     X = Tensor.randn(10, 10).realize()
     idxs = Tensor([0, 2]).realize()
     xt = X[idxs]
-    run_schedule(check_schedule(xt, 1))
+    run_linear(*check_schedule(xt, 1))
     np.testing.assert_equal(xt.numpy(), X.numpy()[idxs.numpy()])
 
   def test_simple_indexing_alt(self):
     X = Tensor.arange(16).reshape(4, 4)
     xt = X[[1, 2], [-1, 2]]
-    run_schedule(check_schedule(xt, 1))
+    run_linear(*check_schedule(xt, 1))
     np.testing.assert_equal(xt.numpy(), (np.arange(16).reshape(4, 4))[[1, 2], [-1, 2]])
 
   def test_advanced_indexing(self):
     X = Tensor.arange(10)+1
     xt = X[[0, -1]]
-    run_schedule(check_schedule(xt, 1))
+    run_linear(*check_schedule(xt, 1))
     np.testing.assert_equal(xt.numpy(), (np.arange(10)+1)[[0, -1]])
 
   def test_advanced_indexing_alt(self):
     X = Tensor.arange(6).reshape(3, 2)+1
     xt = X[[Tensor([2]), Tensor([1])]]
-    run_schedule(check_schedule(xt, 1))
+    run_linear(*check_schedule(xt, 1))
     np.testing.assert_equal(xt.numpy(), 6)
 
   def test_push_through_reshape(self):
     Tensor.manual_seed(0)
     x = Tensor.randn(10, 20).realize()
     out = x.argmax(1)
-    run_schedule(check_schedule(out, 2))
+    run_linear(*check_schedule(out, 2))
     np.testing.assert_allclose(out.numpy(), np.argmax(x.numpy(), 1))
 
   def test_arange_push_through_expand(self):
@@ -933,35 +935,35 @@ class TestSchedule(unittest.TestCase):
     a = Tensor.arange(4,)
     b = Tensor.randn(4, 4).realize()
     out = (a+b).sum()
-    run_schedule(check_schedule(out, 1))
+    run_linear(*check_schedule(out, 1))
     np.testing.assert_allclose(out.numpy(), (np.arange(4)+b.numpy()).sum(), atol=1e-5)
 
   def test_argmin(self):
     Tensor.manual_seed(0)
     x = Tensor.randn(4, 32).realize()
     out = x.argmin(-1)
-    run_schedule(check_schedule(out, 2))
+    run_linear(*check_schedule(out, 2))
     np.testing.assert_equal(out.numpy(), x.numpy().argmin(axis=-1))
 
   def test_argmax(self):
     Tensor.manual_seed(0)
     x = Tensor.randn(4, 32).realize()
     out = x.argmax(-1)
-    run_schedule(check_schedule(out, 2))
+    run_linear(*check_schedule(out, 2))
     np.testing.assert_equal(out.numpy(), x.numpy().argmax(axis=-1))
 
   def test_arange_transposed(self):
     Tensor.manual_seed(0)
     x = Tensor.randint(4, 1).realize()
     a = ((Tensor.arange(4,)*x).T).sum()
-    run_schedule(check_schedule(a, 1))
+    run_linear(*check_schedule(a, 1))
     np.testing.assert_equal(a.numpy(), (np.arange(4)*x.numpy()).T.sum())
 
   def test_div_padded_arange(self):
     x = Tensor.full((2,2), 16)
     y = x.idiv(Tensor.linspace(2, 8, steps=4, dtype=dtypes.int).reshape(2,2)).pad(((1,1), (1,1)))
     out = y.sum(axis=1)
-    run_schedule(check_schedule(out, 1))
+    run_linear(*check_schedule(out, 1))
     self.assertListEqual(out.tolist(), [0, 12, 4, 0])
 
   def test_arange_transposed_descendants(self):
@@ -970,7 +972,7 @@ class TestSchedule(unittest.TestCase):
     a = (Tensor.arange(4,)*x).T
     b = Tensor.randint(4, 4).realize()
     out = (a+b).sum()
-    run_schedule(check_schedule(out, 1))
+    run_linear(*check_schedule(out, 1))
     np.testing.assert_equal(out.numpy(), ((np.arange(4)*x.numpy()).T+b.numpy()).sum())
 
   def test_arange_index(self):
@@ -978,7 +980,7 @@ class TestSchedule(unittest.TestCase):
     x = Tensor.randn(5, 2).realize()
     a = Tensor.arange(10)
     out = (x + a[2]).sum()
-    run_schedule(check_schedule(out, 1))
+    run_linear(*check_schedule(out, 1))
     np.testing.assert_allclose(out.numpy(), (x.numpy()+np.arange(10)[2]).sum(), atol=1e-5, rtol=1e-6)
 
   def test_arange_index_contiguous(self):
@@ -986,7 +988,7 @@ class TestSchedule(unittest.TestCase):
     x = Tensor.randn(5, 2).realize()
     a = Tensor.arange(10).contiguous()
     out = (x + a[2]).sum()
-    run_schedule(check_schedule(out, 2))
+    run_linear(*check_schedule(out, 2))
     np.testing.assert_allclose(out.numpy(), (x.numpy()+np.arange(10)[2]).sum(), atol=1e-5, rtol=1e-6)
 
   def test_arange_index_child(self):
@@ -994,7 +996,7 @@ class TestSchedule(unittest.TestCase):
     x = Tensor.randn(5, 2).realize()
     a = Tensor.arange(10)+1
     out = (x + a[2]).sum()
-    run_schedule(check_schedule(out, 1))
+    run_linear(*check_schedule(out, 1))
     np.testing.assert_allclose(out.numpy(), (x.numpy()+(np.arange(10)+1)[2]).sum(), atol=1e-5, rtol=1e-6)
 
   def test_user_contiguous(self):
@@ -1002,13 +1004,13 @@ class TestSchedule(unittest.TestCase):
     x = Tensor.randn(5, 2).realize()
     a = (Tensor.arange(10)+1).contiguous()
     out = (x + a[2]).sum()
-    run_schedule(check_schedule(out, 2))
+    run_linear(*check_schedule(out, 2))
     np.testing.assert_allclose(out.numpy(), (x.numpy()+(np.arange(10)+1)[2]).sum(), atol=1e-5, rtol=1e-6)
 
   @unittest.skip("BUFFER_VIEW no longer supported on non-disk devices")
   def test_arange_view_op(self):
     a = Tensor.arange(12).reshape(4, 3).shrink(((1, 2), (1, 3))).contiguous()
-    sched = run_schedule(check_schedule(a, 1))
+    sched = run_linear(*check_schedule(a, 1))
     self.assertIs(sched[1].ast.op, Ops.BUFFER_VIEW)
     np.testing.assert_equal(a.numpy(), [[4, 5]])
 
@@ -1017,16 +1019,16 @@ class TestSchedule(unittest.TestCase):
     from extra.models.llama import precompute_freqs_cis
     args = {"dim":32, "end":2048, "theta":10000}
     fused = precompute_freqs_cis(**args)
-    run_schedule(check_schedule(fused, 1))
+    run_linear(*check_schedule(fused, 1))
     if getenv("CHECK", 1):
       ref = precompute_freqs_cis(**args)
-      run_schedule(check_schedule(ref, 1))
+      run_linear(*check_schedule(ref, 1))
       np.testing.assert_equal(fused.numpy(), ref.numpy())
 
   def test_fuse_assign_contiguous(self):
     x = Tensor.zeros(4, 4, dtype=dtypes.int).contiguous().realize()
     a = Tensor.arange(8).reshape(4, 2)
-    run_schedule(check_schedule(x.shrink((None, (0, 2))).assign(a.contiguous()), 2))
+    run_linear(*check_schedule(x.shrink((None, (0, 2))).assign(a.contiguous()), 2))
     np.testing.assert_equal(x.numpy(), [[0, 1, 0, 0], [2, 3, 0, 0], [4, 5, 0, 0], [6, 7, 0, 0]])
 
   def test_assign_non_contiguous_alt(self): self.test_assign_non_contiguous(alt=True)
@@ -1075,8 +1077,8 @@ class TestSchedule(unittest.TestCase):
     idx = Tensor([1,2,5,6], dtype=dtypes.int32)
     flat_base[idx] = Tensor([99,99,99,99])
     base.assign(flat_base.reshape(4, 4))
-    sched = check_schedule(base, 6)  # TODO: this is high
-    run_schedule(sched)
+    sched = check_schedule(base, 4)
+    run_linear(*sched)
     expected = list(range(16))
     for i, v in zip([1,2,5,6], [99,99,99,99]): expected[i] = v
     np.testing.assert_equal(base.reshape(16).numpy(), expected)
@@ -1085,7 +1087,7 @@ class TestSchedule(unittest.TestCase):
     X = Tensor([[0, 2, 3], [1, 2, 3]]).realize()
     Y = Tensor([1, 2]).realize()
     loss = X.sparse_categorical_crossentropy(Y)
-    run_schedule(check_schedule(loss, 3))
+    run_linear(*check_schedule(loss, 3))
     np.testing.assert_allclose(loss.item(), 0.878309, atol=1e-5, rtol=1e-6)
 
   def test_const_folding_alt(self):
@@ -1105,7 +1107,7 @@ class TestSchedule(unittest.TestCase):
     samples = Tensor.randint(BS:=getenv("BS", 512), high=cast(int,Y_train.shape[-1])).realize()
     yt = Tensor.randn(BS, 10).realize()
     loss = yt.sparse_categorical_crossentropy(Y_train[samples])
-    run_schedule(check_schedule(loss, 4))
+    run_linear(*check_schedule(loss, 4))
     loss_fused = loss.numpy()
     loss_ref = torch.nn.CrossEntropyLoss()(torch.tensor(yt.numpy()), torch.tensor(Y_train.numpy())[torch.tensor(samples.numpy())])
     np.testing.assert_allclose(loss_fused, loss_ref.numpy(), atol=1e-6, rtol=1e-6)
@@ -1115,7 +1117,7 @@ class TestSchedule(unittest.TestCase):
     r = (X+Tensor.arange(16).reshape(4, 4)).sum()
     out0 = r+2
     out1 = r+3
-    run_schedule(check_schedule([out0, out1], 2)) # TODO: 1?
+    run_linear(*check_schedule([out0, out1], 2)) # TODO: 1?
     r_ref = (X.numpy()+np.arange(16).reshape(4, 4)).sum()
     np.testing.assert_allclose(out0.numpy(), r_ref+2, rtol=2e-7)
     np.testing.assert_allclose(out1.numpy(), r_ref+3, rtol=2e-7)
@@ -1130,7 +1132,7 @@ class TestSchedule(unittest.TestCase):
     for shape in [(3, 3), (4, 4)]:
       a = Tensor.ones(*shape).contiguous().realize()
       a.assign(a / 1)
-      run_schedule(check_schedule(a, 0, filter_sink=False))
+      run_linear(*check_schedule(a, 0, filter_sink=False))
       self.assertListEqual(a.tolist(), [[1.]*shape[1]]*shape[0])
 
 class TestLimitBufs(unittest.TestCase):
@@ -1164,7 +1166,7 @@ class TestSwizzle(unittest.TestCase):
       a = Tensor.randint(32, 32).realize()
     r = (a+a).sum(1).sum(0)
     # double reduce collapses to a single reduce
-    run_schedule(check_schedule(r, 1))
+    run_linear(*check_schedule(r, 1))
     self.assertEqual(r.numpy(), (a.numpy()+a.numpy()).sum(1).sum(0))
 
   def test_single_swizzle(self):
@@ -1174,7 +1176,7 @@ class TestSwizzle(unittest.TestCase):
       b = Tensor.ones((1, 1), dtype=a.dtype).contiguous().realize()
     # ADD(REDUCE(RESHAPE(LOAD)), LOAD) to ADD(REDUCE(RESHAPE(LOAD))), RESHAPE(LOAD)
     r = a.sum(0)+b
-    run_schedule(check_schedule(r, 1))
+    run_linear(*check_schedule(r, 1))
     self.assertEqual(r.numpy(), a.numpy().sum(0)+1)
 
   def test_double_swizzle_possible(self):
@@ -1184,7 +1186,7 @@ class TestSwizzle(unittest.TestCase):
       b = Tensor.randint(4,).realize()
     # parallel reduce!
     add = a.sum(0)+b.sum(0)
-    run_schedule(check_schedule(add, 1))
+    run_linear(*check_schedule(add, 1))
     self.assertEqual(add.numpy(), a.numpy().sum(0)+b.numpy().sum(0))
 
   def test_swizzle_reduceop(self):
@@ -1192,7 +1194,7 @@ class TestSwizzle(unittest.TestCase):
     x = Tensor.randn(4,4).realize()
     y = Tensor.randn(4,4,4).realize()
     out = x.reshape(4,4,1).expand(4,4,4).sum(axis=(1,))+y
-    run_schedule(check_schedule(out, 2)) # TODO: 1?
+    run_linear(*check_schedule(out, 2)) # TODO: 1?
     np.testing.assert_allclose(out.numpy(), np.tile(x.numpy().reshape(4,4,1), (1,1,4)).sum(axis=1)+y.numpy())
 
   def test_permute_rewrite(self):
@@ -1200,7 +1202,7 @@ class TestSwizzle(unittest.TestCase):
     y = Tensor.randn(4, 1, 16).realize()
     z = Tensor.randn(4, 4, 1).realize()
     t = (x*y).sum(axis=(0, 2)).reshape(1, 4, 1).permute(0, 2, 1)+z
-    run_schedule(check_schedule(t, 2)) # TODO: 1?
+    run_linear(*check_schedule(t, 2)) # TODO: 1?
     t_np = (x.numpy()*y.numpy()).sum(axis=(0, 2)).reshape(1, 4, 1).transpose(0, 2, 1)+z.numpy()
     np.testing.assert_allclose(t.numpy(), t_np, atol=1e-6, rtol=1e-3)
 
@@ -1211,14 +1213,14 @@ class TestSwizzle(unittest.TestCase):
     a_reduce = a.sum(axis=(2,), keepdim=True).sum(axis=(1,))
     b_reduce = b.sum(axis=(0,))
     t = a_reduce+b_reduce
-    run_schedule(check_schedule(t, 1))
+    run_linear(*check_schedule(t, 1))
 
   def test_parallel_reduce_possible(self):
     Tensor.manual_seed(0)
     x = Tensor.randn(4, 2, 2).realize()
     y = Tensor.randn(4, 2, 2).realize()
     t = x.sum(axis=1)+y.sum(axis=1)
-    run_schedule(check_schedule(t, 1))
+    run_linear(*check_schedule(t, 1))
     np.testing.assert_allclose(t.numpy(), x.numpy().sum(axis=1)+y.numpy().sum(axis=1), atol=1e-6, rtol=1e-3)
 
   # kernels can only have 1 or n in each dim
@@ -1227,14 +1229,14 @@ class TestSwizzle(unittest.TestCase):
     x = Tensor.randn(4, 2, 2).realize()
     y = Tensor.randn(4, 3, 2).realize()
     t = x.sum(axis=1)+y.sum(axis=1)
-    run_schedule(check_schedule(t, 1))
+    run_linear(*check_schedule(t, 1))
     np.testing.assert_allclose(t.numpy(), x.numpy().sum(axis=1)+y.numpy().sum(axis=1), atol=1e-6, rtol=1e-3)
 
   def test_unsafe_pad(self):
     x = Tensor.full((2,2), 1.0).contiguous()
     y = x*x.sum((1,)).reciprocal()
     t = y.pad(((0,1),None))
-    run_schedule(check_schedule(t, 3))
+    run_linear(*check_schedule(t, 3))
     np.testing.assert_equal(t.numpy(), [[0.5, 0.5], [0.5, 0.5], [0., 0.]])
 
 zero_pm = UPat(Ops.CONST, arg=0)
@@ -1245,7 +1247,7 @@ class TestView(unittest.TestCase):
     # all masked out, degrades to const 0
     b = a.pad(((0, 10), None))[10:]
     sched = check_schedule(b.contiguous(), 1)
-    run_schedule(sched)
+    run_linear(*sched)
     np.testing.assert_equal(b.numpy(), 0)
 
   def test_mask_dim_1(self):
@@ -1254,7 +1256,7 @@ class TestView(unittest.TestCase):
     b = a.pad((None, (0, 10)))[:, 10:]
     assert b.shape == (10, 10)
     sched = check_schedule(b.contiguous(), 1)
-    run_schedule(sched)
+    run_linear(*sched)
     np.testing.assert_equal(b.numpy(), 0)
 
   def test_partial_mask(self):
@@ -1263,7 +1265,7 @@ class TestView(unittest.TestCase):
     b = a.pad(((0, 5), None))[5:]
     assert b.shape == (10, 10)
     sched = check_schedule(b.contiguous(), 1)
-    run_schedule(sched)
+    run_linear(*sched)
     np.testing.assert_allclose(b.numpy(), np.pad(a.numpy(), ((0, 5), (0, 0)))[5:])
 
   # a*VIEW(x), where VIEW(x) = 0
@@ -1274,7 +1276,7 @@ class TestView(unittest.TestCase):
     bv = b.pad(((0, 2),))[-2:]
     # this becomes a late a*0
     late_mul = a*bv
-    run_schedule(check_schedule(late_mul, 2))
+    run_linear(*check_schedule(late_mul, 2))
     # the arange doesn't realize
     #self.assertIsNone(b.uop.base.realized)
     # mul doesn't realize
@@ -1297,14 +1299,14 @@ class TestView(unittest.TestCase):
     # NOTE: no longer checked
     # mul still collapses
     #self.assertIs(late_mul.uop.base.op, Ops.CONST)
-    run_schedule(s)
+    run_linear(*s)
     self.assertEqual(other_child.tolist(), [2, 3, 4])
 
 @unittest.skipIf(Device.DEFAULT == "CPU", "tests copy from another device to cpu")
 class TestCopyFolding(unittest.TestCase):
   def test_const_copy_is_free(self):
     b = Tensor(1).to("CPU") * 4
-    run_schedule(check_schedule(b, 1, filter_sink=False))
+    run_linear(*check_schedule(b, 1, filter_sink=False))
     assert b.item() == 4
 
   def test_one_hot_with_copy(self):
@@ -1314,14 +1316,14 @@ class TestCopyFolding(unittest.TestCase):
 
   def test_const_copy_multi(self):
     x = Tensor.ones(1, device="CPU").to_(["CPU", "CPU:1"]) * 2
-    run_schedule(check_schedule(x, 2, filter_sink=False))
+    run_linear(*check_schedule(x, 2, filter_sink=False))
     self.assertEqual(x.item(), 2.0)
 
   def test_late_const_copy_folding(self):
     a = Tensor.arange(3).realize()
     zeros = Tensor.zeros(3).realize()
     b = (a*zeros).to("CPU") + 1
-    run_schedule(check_schedule(b, 1, filter_sink=False))
+    run_linear(*check_schedule(b, 1, filter_sink=False))
     self.assertListEqual(b.tolist(), [1, 1, 1])
     self.assertEqual(b.device, "CPU")
 
@@ -1361,7 +1363,7 @@ class TestCopyFolding(unittest.TestCase):
     a = Tensor.ones(4, 4).contiguous().realize()
     # use copy_to_device to bypass Tensor.to() shortcircuit and force a real same-device COPY in the graph
     a.assign(Tensor(a.uop.copy_to_device(a.device), a.device))
-    run_schedule(check_schedule(a, 2, filter_sink=False))
+    run_linear(*check_schedule(a, 2, filter_sink=False))
     self.assertListEqual(a.tolist(), [[1.]*4]*4)
 
   def test_clone(self):
@@ -1372,18 +1374,18 @@ class TestCopyFolding(unittest.TestCase):
     a = Tensor.arange(4)
     view = a.shrink(((0, 2),))
     b = view.clone()
-    run_schedule(check_schedule(b, 1, filter_sink=False))
+    run_linear(*check_schedule(b, 1, filter_sink=False))
     self.assertEqual(b.uop.base.buffer.size, 2)
-    self.assertEqual(b.uop.size, 2)
+    self.assertEqual(b.uop.numel(), 2)
     self.assertListEqual(b.tolist(), [0, 1])
 
   def test_expanded_copy(self):
     a = Tensor.arange(2)
     view = a.reshape(2, 1).expand(2, 2)
     b = view.clone()
-    run_schedule(check_schedule(b, 1, filter_sink=False))
+    run_linear(*check_schedule(b, 1, filter_sink=False))
     self.assertEqual(b.uop.base.buffer.size, 4)
-    self.assertEqual(b.uop.size, 4)
+    self.assertEqual(b.uop.numel(), 4)
     self.assertListEqual(b.tolist(), [[0, 0], [1, 1]])
 
   def test_permuted_copy(self):

test/backend/test_subbuffer.py

@@ -1,7 +1,7 @@
 import unittest
 from tinygrad import Device, dtypes, Tensor
 from tinygrad.device import Buffer
-from tinygrad.helpers import Context, getenv
+from tinygrad.helpers import Context, DEV
 from test.helpers import needs_second_gpu
 
 @unittest.skipUnless(hasattr(Device[Device.DEFAULT].allocator, "_offset"), "subbuffer not supported")
@@ -42,7 +42,7 @@ class TestSubBuffer(unittest.TestCase):
     assert out == [102, 103]
 
   @needs_second_gpu
-  @unittest.skipIf(Device.DEFAULT not in {"CUDA", "NV", "AMD"} or getenv("MOCKGPU"), "only NV, AMD, CUDA")
+  @unittest.skipIf(Device.DEFAULT not in {"CUDA", "NV", "AMD"} or DEV.interface.startswith("MOCK"), "only NV, AMD, CUDA")
   def test_subbuffer_transfer(self):
     t = Tensor.arange(0, 10, dtype=dtypes.uint8).realize()
     vt = t[2:5].contiguous().realize()

test/backend/test_transcendental.py

@@ -14,7 +14,7 @@ settings.load_profile("my_profile")
 
 class TestTranscendentalMath(unittest.TestCase):
   @unittest.skipUnless(is_dtype_supported(dtypes.float64), f"no float64 on {Device.DEFAULT}")
-  @unittest.skipIf(getenv("MOCKGPU") and Device.DEFAULT in {"NV", "CUDA"}, "crashed")
+  @unittest.skipIf(DEV.interface.startswith("MOCK") and Device.DEFAULT in {"NV", "CUDA"}, "crashed")
   @given(ht.float64, strat.sampled_from([(Tensor.exp, np.exp), (Tensor.log, np.log), (Tensor.sin, np.sin)]))
   def test_float64(self, x, op):
     if op[0] == Tensor.sin:
@@ -25,7 +25,7 @@ class TestTranscendentalMath(unittest.TestCase):
                                  op[1](np.array([x], dtype=_to_np_dtype(dtypes.float64))),
                                  atol=3e-2, rtol=1e-5)  # sin can have bigger atol for very big x
 
-  @unittest.skipIf(getenv("MOCKGPU") and Device.DEFAULT in {"NV", "CUDA"}, "crashed")
+  @unittest.skipIf(DEV.interface.startswith("MOCK") and Device.DEFAULT in {"NV", "CUDA"}, "crashed")
   @given(ht.float32, strat.sampled_from([(Tensor.exp, np.exp),(Tensor.log, np.log)] +
     ([(Tensor.sin, np.sin)] if is_dtype_supported(dtypes.ulong) else [])))
   def test_float32(self, x, op):
@@ -66,7 +66,7 @@ class TestFromFuzzer(unittest.TestCase):
     if not is_dtype_supported(dtype): return
     if dtype == dtypes.float64:
       # crashes in CI CUDA
-      if getenv("MOCKGPU") and Device.DEFAULT in {"NV", "CUDA"}: return
+      if DEV.interface.startswith("MOCK") and Device.DEFAULT in {"NV", "CUDA"}: return
     def _test_value(n: float, unit: float=1.0):
       next_float = np.nextafter(1.0, 2.0, dtype=_to_np_dtype(dtype))
       ulp = next_float - 1.0
@@ -88,7 +88,7 @@ class TestFromFuzzer(unittest.TestCase):
     if not is_dtype_supported(dtype): return
     if dtype == dtypes.float64:
       # crashes in CI CUDA
-      if getenv("MOCKGPU") and Device.DEFAULT in {"NV", "CUDA"}: return
+      if DEV.interface.startswith("MOCK") and Device.DEFAULT in {"NV", "CUDA"}: return
     def _test_value(n: float, unit: float=1.0):
       next_float = np.nextafter(1.0, 2.0, dtype=_to_np_dtype(dtype))
       ulp = next_float - 1.0

test/device/test_hcq.py

@@ -1,7 +1,7 @@
 import unittest, ctypes, struct, os, random, numpy as np, time
 from tinygrad import Device, Tensor, dtypes
-from tinygrad.helpers import getenv, mv_address, DEBUG, DEV
-from test.helpers import slow
+from tinygrad.helpers import mv_address, DEBUG, DEV
+from test.helpers import slow, replace_opts
 from tinygrad.device import Buffer, BufferSpec
 from tinygrad.runtime.support.hcq import HCQCompiled, HCQBuffer
 from tinygrad.runtime.autogen import libc
@@ -10,7 +10,7 @@ from tinygrad.engine.realize import get_runner, CompiledRunner, get_program
 from tinygrad.codegen.opt import Opt, OptOps
 from tinygrad import Variable
 
-MOCKGPU = getenv("MOCKGPU")
+MOCKGPU = DEV.interface.startswith("MOCK")
 
 @unittest.skipUnless(issubclass(type(Device[Device.DEFAULT]), HCQCompiled), "HCQ device required to run")
 class TestHCQ(unittest.TestCase):
@@ -76,7 +76,7 @@ class TestHCQ(unittest.TestCase):
         TestHCQ.d0.timeline_signal.wait(TestHCQ.d0.timeline_value)
         TestHCQ.d0.timeline_value += 1
 
-  @unittest.skipIf(Device.DEFAULT in {"CPU"} or (DEV.interface == "PCI" and DEV.device == "AMD"), "Can't handle async update on CPU/MOCKAM device")
+  @unittest.skipIf(Device.DEFAULT == "CPU" or (DEV.interface == "MOCKPCI" and DEV.device == "AMD"), "Can't handle async update on CPU/MOCKPCI device")
   def test_wait_late_set(self):
     for queue_type in [TestHCQ.d0.hw_compute_queue_t, TestHCQ.d0.hw_copy_queue_t]:
       if queue_type is None: continue
@@ -165,7 +165,7 @@ class TestHCQ(unittest.TestCase):
     b = a + 1
     si = b.schedule()[-1]
 
-    runner = CompiledRunner(get_program(si.ast, TestHCQ.d0.renderer, opts=[Opt(op=OptOps.LOCAL, axis=0, arg=3) for _ in range(3)]))
+    runner = CompiledRunner(get_program(replace_opts(si.ast, [Opt(op=OptOps.LOCAL, axis=0, arg=3) for _ in range(3)]), TestHCQ.d0.renderer))
 
     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()
@@ -575,7 +575,7 @@ class TestHCQ(unittest.TestCase):
 
       np.testing.assert_equal(cpu_buffer.numpy(), local_buf.numpy(), "failed")
 
-  @unittest.skipUnless(MOCKGPU and not (DEV.device == "AMD" and DEV.interface == "PCI"), "Emulate this on MOCKGPU to check the path in CI")
+  @unittest.skipUnless(MOCKGPU and not (DEV.device == "AMD" and DEV.interface == "MOCKPCI"), "Emulate this on MOCKGPU to check the path in CI")
   def test_on_device_hang(self):
     if not hasattr(self.d0, 'on_device_hang'): self.skipTest("device does not have on_device_hang")
 

test/external/external_llm_eval.py

@@ -1,4 +1,4 @@
-# eval for tinygrad.apps.llm -- hits the server via OpenAI API
+# eval for OpenAI API server
 # uses Meta's exact ARC-Challenge prompt template from lm-evaluation-harness llama3 tasks
 import argparse, re, pyarrow.parquet as pq
 from openai import OpenAI

test/external/external_test_am_fault_recovery.py

@@ -1,5 +1,6 @@
 # ruff: noqa: F405
 import unittest, subprocess, os
+from tinygrad.helpers import DEV
 from tinygrad.runtime.autogen.amd.rdna3.ins import *  # noqa: F403
 from tinygrad.renderer.amd.dsl import s, v, Inst, NULL
 
@@ -27,7 +28,7 @@ _ILLEGAL_INST_ASM = ".text\n.globl test\n.p2align 8\n.type test,@function\ntest:
   ".rodata\n.p2align 6\n.amdhsa_kernel test\n.amdhsa_next_free_vgpr 8\n.amdhsa_next_free_sgpr 8\n" \
   ".amdhsa_wavefront_size32 1\n.amdhsa_user_sgpr_kernarg_segment_ptr 1\n.amdhsa_kernarg_size 8\n.end_amdhsa_kernel"
 
-@unittest.skipIf(os.environ.get("AMD") != "1" or os.environ.get("MOCKGPU") == "1", "AMD with AM driver required")
+@unittest.skipIf(DEV.device != "AMD" or not DEV.interface.startswith("MOCK"), "AMD with AM driver required")
 class TestAMFaultRecovery(unittest.TestCase):
   def _run_kernel(self, insts: list[Inst]) -> subprocess.CompletedProcess: return _run_asm(assemble_kernel(insts))
 

test/external/external_test_simple_tokenizer.py

@@ -1,7 +1,7 @@
 import functools, multiprocessing
 from transformers import AutoTokenizer
 from datasets import load_dataset
-from tinygrad.apps.llm import SimpleTokenizer
+from tinygrad.llm.cli import SimpleTokenizer
 from tinygrad.helpers import tqdm, getenv, partition
 
 @functools.cache

test/external/process_replay/process_replay.py

@@ -10,9 +10,8 @@ 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.uop.ops import UOp, Ops, KernelInfo
-  from tinygrad.codegen.opt import Opt
-  from tinygrad.helpers import VERSION, Context, ContextVar, colored, db_connection, getenv, tqdm, BEAM
+  from tinygrad.uop.ops import UOp, Ops
+  from tinygrad.helpers import VERSION, Context, ContextVar, colored, db_connection, getenv, tqdm
 except ImportError as e:
   print(repr(e))
   exit(int(ASSERT_DIFF))
@@ -42,13 +41,13 @@ class ProcessReplayWarning(Warning): pass
 
 # *** replay the function and convert return values to string
 
-def replay_get_program(p:ProgramSpec, ast:UOp, renderer:Renderer, opts:list[Opt]|None=None) -> tuple[str, str, tuple[Any, ...]]:
-  if ast.op is Ops.BEAM: ast = ast.src[0]
-  # the ast.arg is non None if we are inside of search.py
-  sink_arg = ast.arg or KernelInfo()
-  if opts is not None: sink_arg = replace(sink_arg, opts_to_apply=tuple(opts))
-  elif BEAM >= 1 and sink_arg.opts_to_apply is None: sink_arg = replace(sink_arg, opts_to_apply=p.applied_opts)
-  input_ast = ast if ast.op is Ops.PROGRAM else ast.replace(arg=replace(sink_arg, name=p.name))
+def replay_get_program(p:ProgramSpec, 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:
     # PYTHON renderer pickles UOps, first unpickle and decode here

test/helpers.py

@@ -1,10 +1,11 @@
 import os, time, struct, functools, unittest
+from dataclasses import replace
 from typing import Any, Callable
 import numpy as np
 from tinygrad import Tensor, dtypes, Device
 from tinygrad.uop.ops import UOp, Ops, KernelInfo
 from tinygrad.tensor import _to_np_dtype
-from tinygrad.engine.realize import Runner, get_program
+from tinygrad.engine.realize import get_program
 from tinygrad.dtype import DType
 from tinygrad.nn.state import get_parameters
 from tinygrad.helpers import T, CI, Target
@@ -23,23 +24,36 @@ def get_uops(sink:UOp, ren:Renderer|None=None) -> list[UOp]:
   full_sink = full_rewrite_to_sink(sink, ren, optimize=sink.tag is None)
   return line_rewrite(linearize(full_sink), pm_linearize_cleanups)
 
+def replace_opts(ast:UOp, opts:list) -> UOp: return ast.replace(arg=replace(ast.arg, opts_to_apply=tuple(opts)))
+
 def derandomize_model(model):
   for p in get_parameters(model):
     p.replace(Tensor.empty(p.shape, device=p.device, dtype=p.dtype))
     p.realize()
 
+def call_is_graph(call:UOp) -> bool:
+  ast = call.src[0]
+  return ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "graph"
+
+def jit_cache_count(linear:UOp) -> int:
+  n = 0
+  for call in linear.src:
+    ast = call.src[0]
+    if ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "graph": n += jit_cache_count(ast.src[0])
+    else: n += 1
+  return n
+
 def assert_jit_cache_len(fxn, expected_len):
-  if not fxn.jit_cache:
+  linear = fxn.captured.linear if fxn.captured is not None else None
+  if linear is None or not linear.src:
     assert expected_len == 0, expected_len
     return
-  # until we have a better way of typing the prg in ExecItem
-  if issubclass(type(fxn.jit_cache[0].prg), Runner) and not type(fxn.jit_cache[0].prg).__name__.endswith('Graph'):
-    assert len(fxn.jit_cache) == expected_len, f"expected {expected_len}, got {len(fxn.jit_cache)}"
+  if call_is_graph(linear.src[0]):
+    assert len(linear.src) == 1, len(linear.src)
+    inner = linear.src[0].src[0].src[0]  # LINEAR UOp inside CUSTOM_FUNCTION
+    assert len(inner.src) == expected_len, f"expected {expected_len}, got {len(inner.src)}"
   else:
-    assert len(fxn.jit_cache) == 1, len(fxn.jit_cache)
-    # until we have a better way of typing the prg in ExecItem
-    assert type(fxn.jit_cache[0].prg).__name__.endswith('Graph')
-    assert len(fxn.jit_cache[0].prg.jit_cache) == expected_len, f"expected {expected_len}, got {len(fxn.jit_cache[0].prg.jit_cache)}"
+    assert len(linear.src) == expected_len, f"expected {expected_len}, got {len(linear.src)}"
 
 def rand_for_dtype(dt:DType, size:int, allow_subnormal=True):
   if dtypes.is_unsigned(dt):

test/mockgpu/amd/README

@@ -4,7 +4,7 @@ Test with `pytest -n12 test/amd/`
 `DEV=AMD:LLVM pytest -n12 test/amd/`
 
 * dsl.py -- helpers for the autogen instruction classes in `__init__.py`. should be standalone with init
-* test/mockgpu/amd/emu.py -- an emulator for RDNA that runs in tinygrad with `DEV=AMD MOCKGPU=1`
+* test/mockgpu/amd/emu.py -- an emulator for RDNA that runs in tinygrad with `DEV=MOCK{KFD|KFD|USB}+AMD`
 * generate.py -- extract assembly format + instruction pseudocode from AMD XML + PDF
 * test/mockgpu/amd/pcode.py -- pseudocode to UOp transformation
 * sqtt.py -- SQTT parser
@@ -20,18 +20,18 @@ 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=AMD MOCKGPU=1 pytest -n=12 test/backend/test_dtype_alu.py test/backend/test_dtype.py`
-`DEV=AMD:LLVM MOCKGPU=1 pytest -n=12 test/backend/test_dtype_alu.py test/backend/test_dtype.py`
+`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`
 
 The ops tests also pass, but they are very slow, so you should run them one at a time.
 
-`SKIP_SLOW_TEST=1 DEV=AMD MOCKGPU=1 pytest -n=12 test/backend/test_ops.py`
-`SKIP_SLOW_TEST=1 DEV=AMD:LLVM MOCKGPU=1 pytest -n=12 test/backend/test_ops.py`
+`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`
 
 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=AMD MOCKGPU=1` it's because an instruction is emulated incorrectly.
-You can test without `MOCKGPU=1` to test on real hardware, if it works on real hardware there's a bug in the emulator.
+If a test is failing with `DEV=MOCKKFD+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.
 
 Currently, only RDNA3 is well supported, but when finished, this will support RDNA3+RDNA4+CDNA in ~3000 lines.

test/mockgpu/amd/amdgpu.py

@@ -7,7 +7,7 @@ from tinygrad.runtime.support import c
 
 MOCKGPU_ARCH = "cdna4" if DEV.arch == "gfx950" else "rdna4" if DEV.arch.startswith("gfx12") else "rdna3"
 assert (ma:=getenv("MOCKGPU_ARCH", "")) == "", "MOCKGPU_ARCH is deprecated, use DEV=" + \
-  str(replace(DEV.value, arch={"cdna4":"gfx950", "rdna4":"gfx1201"}.get(ma, "gfx1100")))
+  str(replace(DEV.value, arch={"cdna4":"gfx950", "rdna4":"gfx1201"}.get(ma, "gfx1100"))) # type: ignore
 GFX_TARGET_VERSION = {"rdna3": 110000, "rdna4": 120000, "cdna4": 90500}[MOCKGPU_ARCH]
 import tinygrad.runtime.autogen.amd_gpu as amd_gpu, tinygrad.runtime.autogen.am.pm4_nv as pm4
 

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
 
-_amd_iface = DEV.target("AMD").interface
-drivers = [NVDriver(), AMDriver() if _amd_iface == "PCI" else (AMUSBDriver() if _amd_iface == "USB" else AMDDriver())]
+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}"))]
 tracked_fds = {}
 
 original_memoryview = builtins.memoryview

test/null/test_attention.py

@@ -1,6 +1,6 @@
 import unittest
 from tinygrad import Tensor, dtypes, TinyJit, UOp
-from tinygrad.apps.llm import apply_rope as apply_rope_new, precompute_freqs_cis
+from tinygrad.llm.model import apply_rope as apply_rope_new, precompute_freqs_cis
 from test.helpers import assert_jit_cache_len
 
 def apply_rope(x:Tensor, start_pos:int):

test/null/test_elf.py

@@ -23,7 +23,7 @@ class TestElfLoader(unittest.TestCase):
       }
     '''
     with self.assertRaisesRegex(RuntimeError, 'evil_external_function'):
-      ClangJITCompiler().compile(src)
+      ClangJITCompiler({'AMD64':'x86_64', 'aarch64':'arm64'}.get(m:=platform.machine(), m)+",native").compile(src)
   def test_link(self):
     src = '''
       float powf(float, float); // from libm

test/null/test_llm_chat.py

@@ -0,0 +1,100 @@
+import unittest
+from tinygrad.llm.cli import SimpleTokenizer, Chat
+
+class TestChatSimple(unittest.TestCase):
+  def test_tekken_preset(self):
+    # Tekken (Mistral): role(user)=[INST], role(assistant)=[], end_turn=[/INST].
+    kv = {
+      "tokenizer.ggml.tokens": ["<unk>", "<s>", "</s>", "[INST]", "[/INST]", "hello", "sure"],
+      "tokenizer.ggml.token_type": [3, 3, 3, 3, 3, 1, 1],
+      "tokenizer.ggml.pre": "tekken",
+      "tokenizer.ggml.eos_token_id": 2,
+    }
+    tok = SimpleTokenizer.from_gguf_kv(kv)
+    chat = Chat.from_gguf_kv(kv, tok)
+    # single user turn: [INST] hello [/INST]
+    self.assertEqual(chat.apply([{"role": "user", "content": "hello"}]), [3, 5, 4])
+    # user + assistant: [INST] hello [/INST] sure [/INST]
+    self.assertEqual(chat.apply([{"role": "user", "content": "hello"}, {"role": "assistant", "content": "sure"}]),
+                     [3, 5, 4, 6, 4])
+    # add_generation_prompt on tekken appends role("assistant") which is []
+    self.assertEqual(chat.apply([{"role": "user", "content": "hello"}], add_generation_prompt=True), [3, 5, 4])
+
+  def test_is_end_basic(self):
+    kv = {"tokenizer.ggml.tokens": ["<unk>", "<eos>"], "tokenizer.ggml.token_type": [3, 3],
+          "tokenizer.ggml.pre": "llama3", "tokenizer.ggml.eos_token_id": 1}
+    tok = SimpleTokenizer.from_gguf_kv(kv)
+    chat = Chat(tok, preset="llama3")
+    self.assertTrue(chat.is_end(1))
+    self.assertFalse(chat.is_end(0))
+
+  def test_olmo2_simple_mode(self):
+    # OLMo 2: pre='dbrx' (unsupported) but arch override maps it to qwen2-style chat with <|im_end|> as turn-end.
+    # tokenizer's eos_id must be untouched; Chat widens stop_ids and uses <|im_end|> for end_turn.
+    kv = {
+      "tokenizer.ggml.tokens": ["<|endoftext|>", "hello", "<|im_end|>", "<|im_start|>"],
+      "tokenizer.ggml.token_type": [3, 1, 3, 3],
+      "tokenizer.ggml.pre": "dbrx",
+      "tokenizer.ggml.eos_token_id": 0,
+      "general.architecture": "olmo2",
+    }
+    tok = SimpleTokenizer.from_gguf_kv(kv)
+    chat = Chat.from_gguf_kv(kv, tok)
+    self.assertEqual(tok.eos_id, 0)                # raw GGUF eos (<|endoftext|>) untouched
+    self.assertEqual(chat.preset, "qwen2")         # arch-overridden preset
+    self.assertEqual(chat.turn_end_id, 2)          # <|im_end|>
+    self.assertTrue(chat.is_end(0))                # <|endoftext|>
+    self.assertTrue(chat.is_end(2))                # <|im_end|>
+    self.assertFalse(chat.is_end(1))
+
+  def test_assistant_prefill_no_end_turn(self):
+    # continue_final_message should drop the trailing end_turn after the prefill assistant message.
+    kv = {
+      "tokenizer.ggml.tokens": ["<unk>", "<s>", "</s>", "[INST]", "[/INST]", "hello", "sure"],
+      "tokenizer.ggml.token_type": [3, 3, 3, 3, 3, 1, 1],
+      "tokenizer.ggml.pre": "tekken",
+      "tokenizer.ggml.eos_token_id": 2,
+    }
+    tok = SimpleTokenizer.from_gguf_kv(kv)
+    chat = Chat.from_gguf_kv(kv, tok)
+    msgs = [{"role": "user", "content": "hello"}, {"role": "assistant", "content": "sure"}]
+    self.assertEqual(chat.apply(msgs), [3, 5, 4, 6, 4])
+    self.assertEqual(chat.apply(msgs, continue_final_message=True), [3, 5, 4, 6])
+
+  def test_unsupported_preset_without_jinja(self):
+    kv = {"tokenizer.ggml.tokens": ["<unk>"], "tokenizer.ggml.token_type": [3], "tokenizer.ggml.pre": "dbrx"}
+    tok = SimpleTokenizer.from_gguf_kv(kv)
+    with self.assertRaises(ValueError):
+      Chat.from_gguf_kv(kv, tok)   # "dbrx" not in _PRESETS and not olmo2 arch
+    # works with use_jinja=True (no simple-preset check), as long as a chat_template is provided
+    kv2 = {**kv, "tokenizer.chat_template": "x"}
+    tok2 = SimpleTokenizer.from_gguf_kv(kv2)
+    Chat.from_gguf_kv(kv2, tok2, use_jinja=True)
+
+class TestChatJinja(unittest.TestCase):
+  def setUp(self):
+    try: import jinja2  # noqa: F401
+    except ImportError: self.skipTest("jinja2 not installed")
+
+  def test_render_simple_template(self):
+    template = ("{%- for m in messages %}{%- if m.role == 'user' %}{{- '[INST]' + m.content + '[/INST]' }}"
+                "{%- elif m.role == 'assistant' %}{{- m.content + '</s>' }}{%- endif %}{%- endfor %}")
+    kv = {
+      "tokenizer.ggml.tokens": ["<unk>", "<s>", "</s>", "[INST]", "[/INST]", "hello"],
+      "tokenizer.ggml.token_type": [3, 3, 3, 3, 3, 1],
+      "tokenizer.ggml.pre": "tekken",
+      "tokenizer.ggml.eos_token_id": 2,
+      "tokenizer.chat_template": template,
+    }
+    tok = SimpleTokenizer.from_gguf_kv(kv)
+    chat = Chat.from_gguf_kv(kv, tok, use_jinja=True)
+    self.assertEqual(chat.apply([{"role": "user", "content": "hello"}]), [3, 5, 4])
+
+  def test_jinja_requires_template(self):
+    kv = {"tokenizer.ggml.tokens": ["<unk>"], "tokenizer.ggml.token_type": [3], "tokenizer.ggml.pre": "llama3"}
+    tok = SimpleTokenizer.from_gguf_kv(kv)
+    with self.assertRaises(ValueError):
+      Chat.from_gguf_kv(kv, tok, use_jinja=True)   # no tokenizer.chat_template
+
+if __name__ == '__main__':
+  unittest.main()

test/null/test_llm_server.py

@@ -7,31 +7,25 @@ class TestLLMServer(unittest.TestCase):
   @classmethod
   def setUpClass(cls):
     cls.mock_tok = Mock()
-    cls.mock_tok.role = Mock(return_value=[100, 101])
     cls.mock_tok.encode = Mock(return_value=[200, 201, 202])
     cls.mock_tok.decode = Mock(return_value="Hello")
     cls.mock_tok.stream_decoder = Mock(return_value=lambda tid=None: "Hello" if tid is not None else "")
-    cls.mock_tok.end_turn = Mock(return_value=[998])
-    cls.mock_tok.preset = "llama3"
+    cls.mock_tok.bos_id = 1
+    cls.mock_tok.eos_id = 999
+    cls.mock_tok.eot_id = None
+
+    cls.mock_chat = Mock()
+    cls.mock_chat.tok = cls.mock_tok
+    cls.mock_chat.apply = Mock(return_value=[1, 100, 101, 200, 201, 202, 998, 100, 101])
+    cls.mock_chat.is_end = Mock(side_effect=lambda tid: tid in (999,))
 
     cls.mock_model = Mock()
     cls.mock_model.generate = Mock(side_effect=lambda ids, **kwargs: iter([300, 301, 999]))
     cls.mock_model.get_start_pos = Mock(return_value=0)
 
-    cls.bos_id = 1
-    cls.eos_id = 999
+    from tinygrad.llm.cli import LLMServer
 
-    import tinygrad.apps.llm as llm_module
-    llm_module.model = cls.mock_model
-    llm_module.model_name = "test-model"
-    llm_module.tok = cls.mock_tok
-    llm_module.bos_id = cls.bos_id
-    llm_module.eos_id = cls.eos_id
-
-    from tinygrad.apps.llm import Handler
-    from tinygrad.viz.serve import TCPServerWithReuse
-
-    cls.server = TCPServerWithReuse(('127.0.0.1', 0), Handler)
+    cls.server = LLMServer(('127.0.0.1', 0), cls.mock_model, "test-model", cls.mock_chat)
     cls.port = cls.server.server_address[1]
     cls.server_thread = threading.Thread(target=cls.server.serve_forever, daemon=True)
     cls.server_thread.start()
@@ -156,35 +150,24 @@ class TestLLMServer(unittest.TestCase):
     self.assertEqual(resp.usage.completion_tokens, 2)
 
   def test_assistant_prefill(self):
-    """Last assistant message should be treated as prefill (not a completed turn)."""
+    """Last assistant message should be treated as prefill (continue_final_message=True)."""
     self.mock_model.generate = Mock(side_effect=lambda ids, **kwargs: iter([300, 999]))
-    captured_ids = []
-    orig_generate = self.mock_model.generate.side_effect
-    def capture_generate(ids, **kwargs):
-      captured_ids.extend(ids)
-      return orig_generate(ids, **kwargs)
-    self.mock_model.generate = Mock(side_effect=capture_generate)
-
+    self.mock_chat.apply.reset_mock()
     resp = self.client.chat.completions.create(
       model="test", messages=[
         {"role": "user", "content": "Hello"},
         {"role": "assistant", "content": "Sure"}
       ], stream=False
     )
-    # prefill tokens should be in ids: role("assistant") + encode("Sure") but NO end_turn after it
-    # and NO extra role("assistant") appended
-    role_tokens = self.mock_tok.role.call_args_list
-    # last role() call should be for "assistant" (the prefill message), not an extra one
-    self.assertEqual(role_tokens[-1], unittest.mock.call("assistant"))
-    # end_turn should be called once less than role() — the prefill assistant msg doesn't get end_turn
-    self.assertEqual(self.mock_tok.end_turn.call_count, self.mock_tok.role.call_count - 1)
+    call = self.mock_chat.apply.call_args
+    self.assertTrue(call.kwargs["continue_final_message"])
+    self.assertFalse(call.kwargs["add_generation_prompt"])
     self.assertIsNotNone(resp.choices[0].message.content)
 
   def test_assistant_prefill_not_last(self):
-    """Assistant message that's NOT last should be a normal completed turn."""
+    """Assistant message that's NOT last should be a normal completed turn (add_generation_prompt=True)."""
     self.mock_model.generate = Mock(side_effect=lambda ids, **kwargs: iter([300, 999]))
-    self.mock_tok.role.reset_mock()
-    self.mock_tok.end_turn.reset_mock()
+    self.mock_chat.apply.reset_mock()
     self.client.chat.completions.create(
       model="test", messages=[
         {"role": "user", "content": "Hello"},
@@ -192,11 +175,9 @@ class TestLLMServer(unittest.TestCase):
         {"role": "user", "content": "Continue"}
       ], stream=False
     )
-    # all messages get end_turn, plus an extra role("assistant") at the end
-    # roles: user, assistant, user, assistant(generation prompt) = 4 role calls
-    # end_turns: user, assistant, user = 3 end_turn calls (one per message)
-    self.assertEqual(self.mock_tok.end_turn.call_count, 3)
-    self.assertEqual(self.mock_tok.role.call_count, 4)
+    call = self.mock_chat.apply.call_args
+    self.assertFalse(call.kwargs["continue_final_message"])
+    self.assertTrue(call.kwargs["add_generation_prompt"])
 
   def test_models_endpoint(self):
     import requests as req

test/null/test_llm_tokenizer.py

@@ -1,5 +1,5 @@
 import unittest, base64, functools, sys
-from tinygrad.apps.llm import SimpleTokenizer
+from tinygrad.llm.cli import SimpleTokenizer
 from tinygrad.helpers import fetch
 
 @unittest.skipIf(sys.platform == 'win32', "fetch race condition on Windows")
@@ -46,18 +46,6 @@ class TestLLMTokenizer(unittest.TestCase):
   def test_llama_repeat(self): self._test_coding(self.llama_tok, "00000000000000000", [ 931, 931, 931, 931, 931, 410 ])
   def test_llama_pat(self): self._test_coding(self.llama_tok, "today\n  \n", [ 31213, 14211 ])
 
-  def test_tekken_from_gguf_kv(self):
-    kv = {
-      "tokenizer.ggml.tokens": ["<unk>", "<s>", "</s>", "[INST]", "[/INST]", "hello"],
-      "tokenizer.ggml.token_type": [3, 3, 3, 3, 3, 1],
-      "tokenizer.ggml.pre": "tekken",
-    }
-    tok = SimpleTokenizer.from_gguf_kv(kv)
-    self.assertEqual(tok.role("user"), [3])
-    self.assertEqual(tok.encode("hello"), [5])
-    self.assertEqual(tok.end_turn(2), [4])
-    self.assertEqual(tok.role("assistant"), [])
-
   def test_stream_decoder(self):
     """stream_decoder buffers incomplete UTF-8: token 25677 has 3/4 of emoji, token 138 completes it."""
     bs = [*range(33, 127), *range(161, 173), *range(174, 256)]

test/null/test_multitensor.py

@@ -153,6 +153,44 @@ class TestMultiRamUsage(unittest.TestCase):
     # with ALLREDUCE_CAST, allreduce copies happen in bf16 (2 bytes) instead of fp32 (4 bytes)
     self.assertLess(mem[1], mem[0])
 
+class TestMultiScalarALU(unittest.TestCase):
+  """Test that tuple-device scalars work correctly in ALU with MULTI tensors (_shard scalar fix)."""
+  def test_multi_times_replicated_scalar(self):
+    devices = ("NULL:0", "NULL:1")
+    x = Tensor.ones(4).contiguous().shard(devices, axis=0)
+    s = Tensor(2.0).to(devices)
+    result = x * s
+    self.assertEqual(result.shape, (4,))
+    self.assertEqual(result.uop.axis, 0)
+
+  def test_multi_add_replicated_scalar(self):
+    devices = ("NULL:0", "NULL:1")
+    x = Tensor.ones(4).contiguous().shard(devices, axis=0)
+    s = Tensor(1.0).to(devices)
+    result = x + s
+    self.assertEqual(result.shape, (4,))
+    self.assertEqual(result.uop.axis, 0)
+
+  def test_multi_times_call_scalar(self):
+    """Per-device scalar from a CALL (like FP8 local amax) used in ALU with MULTI."""
+    import functools
+    from tinygrad.uop.ops import Ops
+    devices = ("NULL:0", "NULL:1")
+    x = Tensor.ones(4, 4).contiguous().shard(devices, axis=0)
+    # simulate per-device scalar via CALL (strips MULTI from param body → no allreduce)
+    @functools.cache
+    def _fxn(x_p, device):
+      t = Tensor(x_p, device=device)
+      inner = Tensor(t.uop.src[0]) if t.uop.op is Ops.MULTI else t
+      return (inner.sum(),)
+    param = x.as_param(0)
+    fxn = _fxn(param.uop, x.device)
+    per_dev_scalar = Tensor(fxn[0].uop.call(x.uop).gettuple(0))
+    result = x * per_dev_scalar
+    self.assertEqual(result.shape, (4, 4))
+    self.assertEqual(result.uop.axis, 0)
+    result.realize()
+
 class TestMultiAxis(unittest.TestCase):
   def test_reshape_shard_invalid(self):
     devices = ("NULL:0", "NULL:1")

test/null/test_process_replay.py

@@ -3,6 +3,7 @@ from tinygrad import Tensor, Device, Context
 from tinygrad.engine.realize import get_program
 from tinygrad.codegen.opt import Opt, OptOps
 from test.external.process_replay.process_replay import replay_get_program
+from test.helpers import replace_opts
 
 N = 16
 class TestProcessReplay(unittest.TestCase):
@@ -19,22 +20,24 @@ class TestProcessReplay(unittest.TestCase):
 
   def test_replay_empty_opts(self):
     # opts=[] means explicitly apply zero opts (unoptimized)
-    p = get_program(self.ast, self.renderer, opts=[])
-    good, compare, _ = replay_get_program(p, self.ast, self.renderer, opts=[])
+    ast = replace_opts(self.ast, [])
+    p = get_program(ast, self.renderer)
+    good, compare, _ = replay_get_program(p, ast, self.renderer)
     self.assertEqual(good, compare)
 
   def test_replay_with_opt(self):
     # opts=[Opt(...)] means apply a specific opt
     opts = [Opt(OptOps.UPCAST, 0, 4)]
-    p = get_program(self.ast, self.renderer, opts=opts)
-    good, compare, _ = replay_get_program(p, self.ast, self.renderer, opts=opts)
+    ast = replace_opts(self.ast, opts)
+    p = get_program(ast, self.renderer)
+    good, compare, _ = replay_get_program(p, ast, self.renderer)
     self.assertEqual(good, compare)
 
-  @Context(BEAM=1)
   def test_beam(self):
-    si = (Tensor.empty(N, N) @ Tensor.empty(N, N)).schedule()[-1]
+    with Context(BEAM=1):
+      si = (Tensor.empty(N, N) @ Tensor.empty(N, N)).schedule()[-1]
     p = get_program(si.ast, self.renderer)
-    good, compare, _ = replay_get_program(p, self.ast, self.renderer)
+    good, compare, _ = replay_get_program(p, si.ast, self.renderer)
     self.assertEqual(good, compare)
 
 if __name__ == '__main__':

test/null/test_real_world.py

@@ -6,7 +6,7 @@ from tinygrad.nn.state import get_parameters
 from tinygrad.engine.jit import TinyJit
 from tinygrad import Tensor, Device, GlobalCounters, dtypes, Variable
 from tinygrad.helpers import Context
-from test.helpers import slow
+from test.helpers import slow, jit_cache_count
 from extra.lr_scheduler import OneCycleLR
 from test.helpers import derandomize_model
 
@@ -31,8 +31,7 @@ def helper_test(nm, gen, model, max_memory_allowed, max_kernels_allowed, all_jit
       tms.append(time.perf_counter_ns() - st)
     mem_used = (GlobalCounters.mem_used - global_mem_used) / 1e9
 
-    # TODO: jit should expose this correctly with graph
-    kernels_used = len(model.jit_cache) if hasattr(model, "jit_cache") else None
+    kernels_used = jit_cache_count(model.captured.linear) if getattr(model, "captured", None) is not None else None
     print(f"{nm}: used {mem_used/1e9:.2f} GB and {kernels_used} kernels in {min(tms)/1e6:.2f} ms")
     assert mem_used < max_memory_allowed, f"{nm} used more than {max_memory_allowed:.3f} GB - {mem_used:.3} GB used"
     assert (max_memory_allowed - mem_used) / max_memory_allowed < 0.2, f"{max_memory_allowed:.3f} GB is too far from {mem_used:.3} GB used"

test/null/test_schedule.py

@@ -3,18 +3,20 @@ import gc, unittest, time
 from tinygrad import nn, dtypes, Device, Tensor
 from tinygrad.uop.ops import UOp, Ops, GroupOp, UPat, KernelInfo
 from tinygrad.helpers import DEBUG, GlobalCounters, Context
-from tinygrad.engine.realize import CompiledRunner, run_schedule
+from tinygrad.engine.realize import CompiledRunner, run_schedule, run_linear
+from tinygrad.schedule import linear_to_schedule
 
 class KernelCountException(Exception): pass
 def check_schedule(t:Tensor|list[Tensor]|UOp, allowed:int, to_prerealize:list[Tensor]|None=None, filter_sink=True):
   if to_prerealize:
     with Context(DEBUG=0, TRACK_MATCH_STATS=0): Tensor.realize(*to_prerealize)
-  if isinstance(t, Tensor): sched = t.schedule()
-  elif isinstance(t, list) and isinstance(t[0], Tensor): sched = Tensor.schedule(*t)
+  if isinstance(t, Tensor): linear, var_vals = t.linear_with_vars()
+  elif isinstance(t, list) and isinstance(t[0], Tensor): linear, var_vals = Tensor.linear_with_vars(*t)
   else:
     assert isinstance(t, UOp), f"can't schedule {t}"
-    sched = Tensor(t).schedule()
+    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])
   if kernel_cnt != allowed:
@@ -24,7 +26,7 @@ def check_schedule(t:Tensor|list[Tensor]|UOp, allowed:int, to_prerealize:list[Te
         print("kernel", i+1)
         print(s.ast)
     raise KernelCountException(f"{kernel_cnt} != {allowed}")
-  return sched
+  return linear, var_vals
 
 def _realize_weights(m):
   for p in nn.state.get_parameters(m): p.realize()
@@ -74,12 +76,12 @@ class TestBufferUOp(unittest.TestCase):
 
   def test_const_does_not_realize(self):
     a = Tensor(1)
-    run_schedule(check_schedule(a, 0))
+    run_linear(*check_schedule(a, 0))
     self.assertIsNone(a.uop.base.realized)
 
   def test_var_does_not_realize(self):
     a = Tensor(UOp.variable("a", 0, 10).bind(1))
-    run_schedule(check_schedule(a, 0))
+    run_linear(*check_schedule(a, 0))
     self.assertIsNone(a.uop.base.realized)
 
   def test_unused_var_not_in_var_vals(self):
@@ -377,8 +379,8 @@ class TestSchedule(unittest.TestCase):
     r0 = x.mean(axis=0, keepdim=True)
     r1 = (x - r0).sum(axis=0).div(2)
     out = r0 + r1
-    schedule = check_schedule(out, 2)
-    reduceops = [x for si in schedule for x in si.ast.toposort() if x.op in {Ops.REDUCE_AXIS, Ops.REDUCE}]
+    linear, _ = check_schedule(out, 2)
+    reduceops = [x for si in linear.src for x in si.src[0].toposort() if x.op in {Ops.REDUCE_AXIS, Ops.REDUCE}]
     assert len(reduceops) == 2
 
   def test_cache_reduce_multiple_children(self):
@@ -388,8 +390,8 @@ class TestSchedule(unittest.TestCase):
     r1 = (x - r0).sum(axis=0).div(2)
     out0 = r0 + y
     out1 = r1 + y
-    schedule = check_schedule([out0, out1], 3)
-    reduceops = [x for si in schedule for x in si.ast.toposort() if x.op in {Ops.REDUCE_AXIS, Ops.REDUCE}]
+    linear, _ = check_schedule([out0, out1], 3)
+    reduceops = [x for si in linear.src for x in si.src[0].toposort() if x.op in {Ops.REDUCE_AXIS, Ops.REDUCE}]
     self.assertEqual(len(reduceops), 2) # why is RANGEIFY different?
 
   def test_dedup_assign(self):
@@ -808,7 +810,7 @@ class TestSchedule(unittest.TestCase):
     Tensor.manual_seed(0)
     x = Tensor.randn(4, 12, 64, 64, requires_grad=True).realize()
     x.softmax().sum().backward()
-    run_schedule(check_schedule(x.grad, 4))
+    run_linear(*check_schedule(x.grad, 4))
 
   def test_scaled_dot_product_attention_fusion(self):
     x, y, z, m = (Tensor.empty(32, 8, 16, 16) for _ in range(4))
@@ -946,7 +948,7 @@ class TestSchedule(unittest.TestCase):
     Tensor.manual_seed(0)
     x = Tensor.randn(10, 20).realize()
     out = x.argmax(1)
-    run_schedule(check_schedule(out, 2))
+    run_linear(*check_schedule(out, 2))
 
   def test_resnet_conv2d(self):
     x = Tensor.empty(1, 8, 32, 32)
@@ -986,15 +988,15 @@ class TestSchedule(unittest.TestCase):
 
   def test_fuse_arange_avg_pool2d_ceil_mode(self):
     x = Tensor.avg_pool2d(Tensor.empty(1,1,6,6), kernel_size=(3,3), padding=1, stride=3, ceil_mode=True)
-    sched = check_schedule(x, 1)
-    self.assertEqual(len([x for x in sched[0].ast.backward_slice_with_self if x.op is Ops.REDUCE]), 1)
+    linear, _ = check_schedule(x, 1)
+    self.assertEqual(len([x for x in linear.src[0].src[0].backward_slice_with_self if x.op is Ops.REDUCE]), 1)
 
   def test_fuse_arange_pad_circular_mode_bw(self):
     x = Tensor.empty(1,1,5,5,5)
     out = x.pad((1,2,3,5,1,2), mode="circular")
     g = out.sum().gradient(x)[0]
-    sched = check_schedule(g, 1)
-    self.assertEqual(len([x for x in sched[0].ast.backward_slice_with_self if x.op is Ops.REDUCE]), 0)
+    linear, _ = check_schedule(g, 1)
+    self.assertEqual(len([x for x in linear.src[0].src[0].backward_slice_with_self if x.op is Ops.REDUCE]), 0)
 
   def test_resnet_block(self):
     with Tensor.train(False):
@@ -1007,7 +1009,7 @@ class TestSchedule(unittest.TestCase):
       out = bn1(conv1(x)).relu()
       out = bn2(conv2(out))
       out = (out + x).relu()
-      run_schedule(check_schedule(out, 2, [conv1.weight, conv2.weight]))
+      run_linear(*check_schedule(out, 2, [conv1.weight, conv2.weight]))
 
 class TestSwizzle(unittest.TestCase):
   def test_softmax_one_kernel(self):
@@ -1041,7 +1043,7 @@ class TestUOpBecome(unittest.TestCase):
     assert UPat(Ops.BUFFER).match(add.uop.base, {})
     # the Tensor UOp can optionally stack a VIEW on top of the BUFFER, in this case to preserve the (4, 4) shape of the tensor
     assert add.uop is not add.uop.base
-    self.assertEqual(add.uop.size, 16)
+    self.assertEqual(add.uop.numel(), 16)
     self.assertEqual(add.uop.shape, (4, 4))
 
   def test_new_buffer_view(self):
@@ -1217,72 +1219,72 @@ class TestBufferView(unittest.TestCase):
     # simple 1D shrink of a realized buffer should be BUFFER_VIEW, not a copy kernel
     a = Tensor.arange(100).contiguous().realize()
     b = a.shrink(((10, 50),)).contiguous()
-    run_schedule(check_schedule(b, 0))
+    run_linear(*check_schedule(b, 0))
 
   def test_shrink_2d_contiguous_is_buffer_view(self):
     a = Tensor.arange(100).reshape(10,10).contiguous().realize()
     b = a.shrink(((1, 5),None)).contiguous()
-    run_schedule(check_schedule(b, 0))
+    run_linear(*check_schedule(b, 0))
 
   def test_chained_shrink_is_buffer_view(self):
     a = Tensor.arange(1000).contiguous().realize()
     b = a.shrink(((200, 800),)).shrink(((0, 300),)).reshape((30, 10)).shrink(((20, 25), (0, 10))).contiguous()
-    run_schedule(check_schedule(b, 0))
+    run_linear(*check_schedule(b, 0))
 
   def test_shrink_non_shard_axis_is_buffer_view_multi(self):
     # indexing a non-shard axis of a realized sharded tensor should be BUFFER_VIEW on each device, not copy kernels
     # this is the flat_llama pattern: weight[layer_idx] where weight is (n_layers, out, dim) sharded on axis=1
     devices = ("NULL:1", "NULL:2")
     a = Tensor.arange(8*4*10).reshape(8, 4, 10).contiguous().shard(devices, axis=1).realize()
-    run_schedule(check_schedule(a[3].contiguous(), 0))
+    run_linear(*check_schedule(a[3].contiguous(), 0))
 
   def test_shrink_2d_non_shard_axis_multi(self):
     devices = ("NULL:1", "NULL:2")
     a = Tensor.arange(6*4).reshape(6, 4).contiguous().shard(devices, axis=1).realize()
-    run_schedule(check_schedule(a.shrink(((1, 4), None)).contiguous(), 0))
+    run_linear(*check_schedule(a.shrink(((1, 4), None)).contiguous(), 0))
 
   def test_shrink_shard_axis_0_multi(self):
     # shrinking a middle dim is not contiguous per shard, so this needs copy kernels
     devices = ("NULL:1", "NULL:2")
     a = Tensor.arange(4*6*2).reshape(4, 6, 2).contiguous().shard(devices, axis=0).realize()
-    run_schedule(check_schedule(a.shrink((None, (2, 5), None)).contiguous(), 2))
+    run_linear(*check_schedule(a.shrink((None, (2, 5), None)).contiguous(), 2))
 
   def test_reshape_then_shrink_multi(self):
     devices = ("NULL:1", "NULL:2")
     a = Tensor.arange(8*6).reshape(8, 6).contiguous().shard(devices, axis=1).realize()
-    run_schedule(check_schedule(a.reshape(4, 2, 6)[1].contiguous(), 0))
+    run_linear(*check_schedule(a.reshape(4, 2, 6)[1].contiguous(), 0))
 
   def test_permute_then_shrink_multi(self):
     # permute makes per-shard view non-contiguous, needs copy kernels
     devices = ("NULL:1", "NULL:2")
     a = Tensor.arange(4*6*2).reshape(4, 6, 2).contiguous().shard(devices, axis=1).realize()
-    run_schedule(check_schedule(a.permute(1, 0, 2).shrink(((0, 6), (1, 3), None)).contiguous(), 2))
+    run_linear(*check_schedule(a.permute(1, 0, 2).shrink(((0, 6), (1, 3), None)).contiguous(), 2))
 
   def test_multi_buffer_view_4_devices(self):
     devices = tuple(f"NULL:{i}" for i in range(4))
     a = Tensor.arange(8*12).reshape(8, 12).contiguous().shard(devices, axis=1).realize()
-    run_schedule(check_schedule(a[5].contiguous(), 0))
+    run_linear(*check_schedule(a[5].contiguous(), 0))
 
   def test_chained_shrink_multi(self):
     devices = ("NULL:1", "NULL:2")
     a = Tensor.arange(10*8).reshape(10, 8).contiguous().shard(devices, axis=1).realize()
-    run_schedule(check_schedule(a.shrink(((2, 8), None)).shrink(((1, 4), None)).contiguous(), 0))
+    run_linear(*check_schedule(a.shrink(((2, 8), None)).shrink(((1, 4), None)).contiguous(), 0))
 
   # negative tests: these should NOT become BUFFER_VIEW (non-contiguous per shard)
   def test_expand_multi_not_buffer_view(self):
     devices = ("NULL:1", "NULL:2")
     a = Tensor.arange(4*2).reshape(4, 1, 2).contiguous().shard(devices, axis=2).realize()
-    run_schedule(check_schedule(a.expand(4, 3, 2).contiguous(), 2))
+    run_linear(*check_schedule(a.expand(4, 3, 2).contiguous(), 2))
 
   def test_pad_multi_not_buffer_view(self):
     devices = ("NULL:1", "NULL:2")
     a = Tensor.arange(4*2).reshape(4, 2).contiguous().shard(devices, axis=1).realize()
-    run_schedule(check_schedule(a.pad(((1, 1), (0, 0))).contiguous(), 2))
+    run_linear(*check_schedule(a.pad(((1, 1), (0, 0))).contiguous(), 2))
 
   def test_flip_multi_not_buffer_view(self):
     devices = ("NULL:1", "NULL:2")
     a = Tensor.arange(4*2).reshape(4, 2).contiguous().shard(devices, axis=1).realize()
-    run_schedule(check_schedule(a.flip(0).contiguous(), 2))
+    run_linear(*check_schedule(a.flip(0).contiguous(), 2))
 
 class TestInvalidTensor(unittest.TestCase):
   def test_full_invalid_is_zero_kernels(self):

test/null/test_tensor.py

@@ -107,10 +107,6 @@ class TestIdxUpcast(unittest.TestCase):
     uops = self._schedule_render(a)
     assert all(uop.dtype is not dtypes.long for uop in uops)
 
-  def test_arange_raise_overflow(self):
-    with self.assertRaises(ValueError):
-      self._schedule_render(Tensor.arange(2**33, dtype=dtypes.int))
-
   @unittest.skipIf(is_dtype_supported(dtypes.long), "int64 is supported")
   def test_int64_unsupported_overflow_sym(self):
     with self.assertRaises((KeyError, RuntimeError)):

test/null/test_tensor_metadata.py

@@ -115,5 +115,14 @@ class TestTensorMetadata(unittest.TestCase):
     c[:4].assign(shared)
     self.assertTrue(self._has_metadata((c[:4] + shared).relu(), "relu"))
 
+class TestTraceMetaShutdown(unittest.TestCase):
+  def test_tracemeta_del_no_shutdown_error(self):
+    import subprocess, os
+    result = subprocess.run(['python3', '-c', 'from tinygrad import Tensor\n'
+                             'x=Tensor.eye(3,requires_grad=True); (x@x).sum().backward()'],
+                            env={**os.environ, "TRACEMETA": "2"}, capture_output=True)
+    self.assertEqual(result.returncode, 0)
+    self.assertNotIn(b"Exception", result.stderr)
+
 if __name__ == '__main__':
   unittest.main()

test/null/test_tensor_uop_mixin.py

@@ -1,5 +1,6 @@
 import math, unittest
-from tinygrad import Tensor
+from tinygrad import Tensor, dtypes
+from tinygrad.uop.ops import UOp
 
 def _t(*shape):
   return Tensor.arange(math.prod(shape)).reshape(*shape)
@@ -70,5 +71,67 @@ class TestTensorUOpStack(unittest.TestCase):
   def test_stack_3tensors(self): _check(self, _t(2, 3), lambda x: x.stack(x, x, dim=0))
   def test_stack_new_last(self): _check(self, _t(2, 3), lambda x: x.stack(x, dim=-1))
 
+class TestTensorUOpConv2d(unittest.TestCase):
+  def test_conv2d_basic(self):
+    w = _t(1, 1, 2, 2).float()
+    _check(self, _t(1, 1, 3, 3).float(), lambda x: x.conv2d(w if isinstance(x, Tensor) else w.uop))
+  def test_conv2d_padded(self):
+    w = _t(1, 1, 2, 2).float()
+    _check(self, _t(1, 1, 3, 3).float(), lambda x: x.conv2d(w if isinstance(x, Tensor) else w.uop, padding=1))
+  def test_conv2d_negative_padding(self):
+    w = _t(1, 1, 3, 3).float()
+    _check(self, _t(1, 1, 5, 5).float(), lambda x: x.conv2d(w if isinstance(x, Tensor) else w.uop, padding=(-1,-1,-1,-1)))
+  def test_conv2d_multichannel_bias(self):
+    w, b = _t(4, 2, 3, 3).float(), _t(4).float()
+    _check(self, _t(2, 2, 5, 5).float(), lambda x: x.conv2d(*(y if isinstance(x, Tensor) else y.uop for y in (w, b))))
+  def test_conv2d_stride_dilation(self):
+    w = _t(2, 2, 2, 2).float()
+    _check(self, _t(1, 2, 6, 6).float(), lambda x: x.conv2d(w if isinstance(x, Tensor) else w.uop, stride=2, dilation=2))
+  def test_conv2d_groups(self):
+    w = _t(4, 1, 2, 2).float()
+    _check(self, _t(1, 4, 4, 4).float(), lambda x: x.conv2d(w if isinstance(x, Tensor) else w.uop, groups=4))
+  def test_conv2d_3d(self):
+    w = _t(1, 1, 2, 2, 2).float()
+    _check(self, _t(1, 1, 3, 3, 3).float(), lambda x: x.conv2d(w if isinstance(x, Tensor) else w.uop))
+  def test_conv_transpose2d_basic(self):
+    w = _t(1, 1, 2, 2).float()
+    _check(self, _t(1, 1, 3, 3).float(), lambda x: x.conv_transpose2d(w if isinstance(x, Tensor) else w.uop))
+  def test_conv_transpose2d_stride(self):
+    w = _t(1, 1, 2, 2).float()
+    _check(self, _t(1, 1, 3, 3).float(), lambda x: x.conv_transpose2d(w if isinstance(x, Tensor) else w.uop, stride=2))
+
+class TestTensorUOpEinsum(unittest.TestCase):
+  def test_einsum_dot(self):       _check(self, _t(2, 3), lambda x: type(x).einsum("ij,ij->", x, x))
+  def test_einsum_transpose(self): _check(self, _t(2, 3), lambda x: type(x).einsum("ij->ji", x))
+
+class TestTensorUOpSoftmax(unittest.TestCase):
+  def test_softmax_default(self):     _check(self, _t(2, 3).float(), lambda x: x.softmax())
+  def test_softmax_axis0(self):       _check(self, _t(2, 3).float(), lambda x: x.softmax(axis=0))
+  def test_log_softmax_default(self): _check(self, _t(2, 3).float(), lambda x: x.log_softmax())
+  def test_log_softmax_axis0(self):   _check(self, _t(2, 3).float(), lambda x: x.log_softmax(axis=0))
+
+# UOp.empty / UOp.empty_like are the canonical buffer allocators; Tensor.empty / Tensor.empty_like just forward.
+class TestUOpEmpty(unittest.TestCase):
+  def test_empty_dtype_string(self):
+    self.assertEqual(UOp.empty((3, 4), dtype="float32").dtype, dtypes.float32)
+
+  def test_empty_like_dtype_override(self):
+    u = Tensor.ones(3, 4).uop.empty_like(dtype=dtypes.int8)
+    self.assertEqual((u.shape, u.dtype), ((3, 4), dtypes.int8))
+    self.assertTrue(u.has_buffer_identity())
+
+  def test_empty_like_sharded_to_single_device(self):
+    # regression: sharded source, override to single device must yield full logical shape with no axis
+    t = Tensor.ones(8, 4).shard(("NULL:0", "NULL:1"), axis=0)
+    for dev in ("NULL:2", ("NULL:2",)):  # singleton tuple also canonicalizes to single device
+      u = t.uop.empty_like(device=dev, dtype=dtypes.int32)
+      self.assertEqual((u.shape, u.device, u.dtype, u.axis), ((8, 4), "NULL:2", dtypes.int32, None))
+      self.assertTrue(u.has_buffer_identity())
+
+  def test_empty_direct_singleton_tuple_device(self):
+    # regression: direct UOp.empty with a singleton-tuple device + axis must not trip .multi()'s tuple assert
+    u = UOp.empty((4,), dtype=dtypes.float32, device=("NULL:0",), axis=0)
+    self.assertEqual((u.shape, u.device, u.axis), ((4,), "NULL", None))
+
 if __name__ == "__main__":
   unittest.main()

test/null/test_uops_stats.py

@@ -9,6 +9,7 @@ from tinygrad.dtype import dtypes
 from tinygrad.codegen.opt import Opt, OptOps, KernelOptError
 from tinygrad.device import Device
 from tinygrad.renderer.ptx import PTXRenderer
+from test.helpers import replace_opts
 
 def flops_mem(uops, ignore_indexing=False):
   est = Estimates.from_uops(uops, ignore_indexing)
@@ -175,13 +176,14 @@ class TestStatsOptimized(unittest.TestCase):
     self.assertEqual(p.estimates.mem, 3*N*N*4) # 3 NxN mats with floats
 
   def test_gemm(self):
-    p = get_program(self.ast_gemm, renderer=Device[Device.DEFAULT].renderer, opts=[])
+    p = get_program(replace_opts(self.ast_gemm, []), renderer=Device[Device.DEFAULT].renderer)
     self.check_gemm(p)
     self.assertEqual(p.estimates.lds, 2*N*N*N*4 + 4*N*N)
 
   def test_gemm_tc_unroll(self):
     try:
-      p = get_program(self.ast_gemm, renderer=Device[Device.DEFAULT].renderer, opts=[Opt(OptOps.TC, 0, (-1, 0, 1)), Opt(OptOps.UNROLL, 0, 2)])
+      p = get_program(replace_opts(self.ast_gemm, [Opt(OptOps.TC, 0, (-1, 0, 1)), Opt(OptOps.UNROLL, 0, 2)]),
+                      renderer=Device[Device.DEFAULT].renderer)
     except KernelOptError:
       raise unittest.SkipTest("no tensor cores")
     print(p.src)
@@ -190,20 +192,20 @@ class TestStatsOptimized(unittest.TestCase):
   # this is a good lesson about why UPCASTing is a good idea
 
   def test_gemm_one_upcasted(self):
-    p = get_program(self.ast_gemm, renderer=Device[Device.DEFAULT].renderer, opts=[Opt(OptOps.UPCAST, 0, 4)])
+    p = get_program(replace_opts(self.ast_gemm, [Opt(OptOps.UPCAST, 0, 4)]), renderer=Device[Device.DEFAULT].renderer)
     self.check_gemm(p)
     self.assertEqual(p.estimates.lds, N*N*N*4 + N*N*N*4//4 + 4*N*N)
 
   def test_gemm_upcasted(self):
-    p = get_program(self.ast_gemm, renderer=Device[Device.DEFAULT].renderer,
-                    opts=[Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UPCAST, 1, 4), Opt(OptOps.UNROLL, 0, 4)])
+    p = get_program(replace_opts(self.ast_gemm, [Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UPCAST, 1, 4), Opt(OptOps.UNROLL, 0, 4)]),
+                    renderer=Device[Device.DEFAULT].renderer)
     self.check_gemm(p)
     self.assertEqual(p.estimates.lds, 2*N*N*N*4//4 + 4*N*N)
 
   def test_gemm_upcasted_locals(self):
     try:
-      p = get_program(self.ast_gemm, renderer=Device[Device.DEFAULT].renderer, opts=[Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UPCAST, 1, 4),
-                                           Opt(OptOps.LOCAL, 0, 4),  Opt(OptOps.LOCAL, 1, 4)])
+      p = get_program(replace_opts(self.ast_gemm, [Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UPCAST, 1, 4), Opt(OptOps.LOCAL, 0, 4),
+                                                   Opt(OptOps.LOCAL, 1, 4)]), renderer=Device[Device.DEFAULT].renderer)
     except KernelOptError:
       raise unittest.SkipTest("no locals")
     self.check_gemm(p)
@@ -211,7 +213,7 @@ class TestStatsOptimized(unittest.TestCase):
 
   def test_gemm_group(self):
     try:
-      p = get_program(self.ast_gemm, renderer=Device[Device.DEFAULT].renderer, opts=[Opt(OptOps.GROUP, 0, 4)])
+      p = get_program(replace_opts(self.ast_gemm, [Opt(OptOps.GROUP, 0, 4)]), renderer=Device[Device.DEFAULT].renderer)
     except KernelOptError:
       raise unittest.SkipTest("no locals")
     SZ = N*N*4
@@ -220,14 +222,14 @@ class TestStatsOptimized(unittest.TestCase):
     self.assertEqual(p.estimates.lds, 2*N*N*N*4 + SZ*4 + (SZ*4 + 4*N*N)*4)
 
   def test_reduce(self):
-    p = get_program(self.ast_reduce, renderer=Device[Device.DEFAULT].renderer, opts=[])
+    p = get_program(replace_opts(self.ast_reduce, []), renderer=Device[Device.DEFAULT].renderer)
     print(p.name, p.estimates.ops, p.estimates.mem, p.estimates.lds)
     self.assertEqual(p.estimates.ops, N*N)
     self.assertEqual(p.estimates.mem, N*N*4 + 4)
 
   def test_reduce_group(self):
     try:
-      p = get_program(self.ast_reduce, renderer=Device[Device.DEFAULT].renderer, opts=[Opt(OptOps.GROUP, 0, 50)])
+      p = get_program(replace_opts(self.ast_reduce, [Opt(OptOps.GROUP, 0, 50)]), renderer=Device[Device.DEFAULT].renderer)
     except KernelOptError:
       raise unittest.SkipTest("no locals")
     # NOTE: these are wrong, they don't respect the if statement

test/null/test_viz.py

@@ -1,12 +1,13 @@
-import unittest, decimal, sys, json, contextlib
+import unittest, decimal, sys, json, contextlib, tempfile, pickle, io, itertools
+from pathlib import Path
 from dataclasses import dataclass
 from typing import Generator
 
-from tinygrad.uop.ops import UOp, UPat, Ops, PatternMatcher, TrackedPatternMatcher, graph_rewrite, track_rewrites, TRACK_MATCH_STATS, profile_matches
+from tinygrad.uop.ops import UOp, UPat, Ops, PatternMatcher, TrackedPatternMatcher, graph_rewrite, track_rewrites, profile_matches
 from tinygrad.uop.symbolic import sym
 from tinygrad.dtype import dtypes
-from tinygrad.helpers import PROFILE, colored, ansistrip, flatten, TracingKey, ProfileRangeEvent, ProfileEvent, Context, cpu_events, profile_marker
-from tinygrad.helpers import VIZ, cpu_profile, ProfilePointEvent, unwrap
+from tinygrad.helpers import colored, ansistrip, flatten, TracingKey, ProfileRangeEvent, ProfileEvent, Context, cpu_events, profile_marker
+from tinygrad.helpers import cpu_profile, ProfilePointEvent, unwrap
 from tinygrad.device import Buffer
 
 from tinygrad.uop.ops import tracked_keys, tracked_ctxs, uop_fields, active_rewrites, active_group, _name_cnt, RewriteTrace
@@ -37,25 +38,13 @@ class VizTrace:
 
 @contextlib.contextmanager
 def save_viz():
-  # clear previous traces
   for lst in [tracked_keys, tracked_ctxs, active_rewrites, active_group, _name_cnt]: lst.clear()
   Buffer.profile_events.clear()
   cpu_events.clear()
-  # set the context vars to enable VIZ
-  prev_viz = VIZ.value
-  VIZ.value = -1
-  prev_tms = TRACK_MATCH_STATS.value
-  TRACK_MATCH_STATS.value = 2
-  prev_profile = PROFILE.value
-  PROFILE.value = 1
   viz = VizTrace()
-  try:
+  with Context(VIZ=-1, TRACK_MATCH_STATS=2, PROFILE=1):
     yield viz
-  finally:
-    viz.set_data()
-    TRACK_MATCH_STATS.value = prev_tms
-    PROFILE.value = prev_profile
-    VIZ.value = prev_viz
+  viz.set_data()
 
 class TestViz(unittest.TestCase):
   def test_simple(self):
@@ -345,7 +334,7 @@ class TestVizIntegration(unittest.TestCase):
       c1 = Tensor.empty(4).add(1)
       c2 = Tensor.empty(8).add(1)
       sched = Tensor.schedule(c1, c2)
-      prgs = [si.lower().prg.p.name for si in sched]
+      prgs = [get_program(si.ast, Device[Device.DEFAULT].renderer).name for si in sched]
     lst = viz.list_items()
     sched_idx = next(i for i,l in enumerate(lst) if l["name"].startswith("Schedule"))
     viz_kernel = next(i for i,s in enumerate(lst[sched_idx]["steps"]) if s["name"] == "View Kernel Graph")
@@ -415,7 +404,7 @@ class TestVizIntegration(unittest.TestCase):
 
 from tinygrad.device import ProfileDeviceEvent, ProfileGraphEvent, ProfileGraphEntry
 from tinygrad.viz.serve import get_profile
-from extra.viz.cli import decode_profile
+from tinygrad.viz.cli import decode_profile
 
 def load_profile(lst:list[ProfileEvent]) -> dict: return decode_profile(get_profile(VizData(), lst))
 
@@ -896,5 +885,47 @@ class TestCfg(unittest.TestCase):
     k.emit(s_code_end())
     self.get_cfg("jump_back_to_end", k)
 
+# launch viz cli without subprocess
+def run_cli(*cli_args) -> str:
+  from tinygrad.viz.cli import main, get_arg_parser
+  args = get_arg_parser().parse_args(cli_args)
+  with contextlib.redirect_stdout(buf:=io.StringIO()):
+    main(args)
+  return buf.getvalue().strip()
+
+class TestCLI(unittest.TestCase):
+  def test_simple(self):
+    a = Tensor.empty(1, device="NULL")+2.0
+    empty_counter = itertools.count(0)
+    def custom_empty_prg(B:UOp, A:UOp) -> UOp:
+      sink = UOp(Ops.SINK, arg=KernelInfo(name=f"custom_empty_n{next(empty_counter)}"))
+      return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=a.device), UOp(Ops.LINEAR, src=(sink,))))
+    b = Tensor.custom_kernel(Tensor.empty_like(a), a, fxn=custom_empty_prg)[0]
+    c = Tensor.custom_kernel(Tensor.empty_like(a), a, fxn=custom_empty_prg)[0]
+    with save_viz() as viz:
+      b.realize()
+      profile_marker("marker @ 1")
+      c.realize()
+    # save trace to disk for CLI to consume it
+    with tempfile.TemporaryDirectory() as tmpdir:
+      (r:=Path(tmpdir)/"rewrites.pkl").write_bytes(pickle.dumps(viz.data.trace))
+      (p:=Path(tmpdir)/"profile.pkl").write_bytes(pickle.dumps(cpu_events))
+      # reconstruct DEBUG=4 output and see all markers.
+      with Context(DEBUG=4):
+        kernels = run_cli("--rewrites-path", str(r), "--profile-path", str(p), "-p", "-s", "NULL")
+      self.assertIn("void custom_empty_n0", kernels)
+      self.assertIn("marker @ 1", kernels)
+      self.assertIn("void custom_empty_n1", kernels)
+      self.assertIn("E", kernels)
+      self.assertIn("UOp.const", kernels)
+      # get the top slowest functions across all devices
+      with Context(DEBUG=2):
+        times = run_cli("--rewrites-path", str(r), "--profile-path", str(p), "-p", "-s", "ALL", "--top", "-1")
+      self.assertIn("TINY", times)
+      self.assertIn("NULL", times)
+      with Context(DEBUG=3):
+        json_lines = run_cli("--rewrites-path", str(r), "--profile-path", str(p), "-p", "-s", "ALL", "--jsonl")
+      for line in json_lines.split("\n"): _ = json.loads(line)
+
 if __name__ == "__main__":
   unittest.main()

test/opt/test_gen_float4.py

@@ -4,6 +4,7 @@ from tinygrad.uop.ops import UOp, Ops
 from tinygrad.codegen.opt import Opt, OptOps
 from tinygrad.engine.realize import get_program
 from tinygrad.helpers import AMX
+from test.helpers import replace_opts
 
 @unittest.skipUnless(Device[Device.DEFAULT].renderer.supports_float4, "need backends that support float4")
 class TestFloat4(unittest.TestCase):
@@ -24,7 +25,7 @@ class TestFloat4(unittest.TestCase):
     s = c.schedule()[0]
     realized_ast = s.ast
     opts_to_apply = [Opt(op=OptOps.UPCAST, axis=0, arg=4)]
-    program = get_program(realized_ast, renderer=Device[Device.DEFAULT].renderer, opts=opts_to_apply)
+    program = get_program(replace_opts(realized_ast, opts_to_apply), renderer=Device[Device.DEFAULT].renderer)
 
     assert TestFloat4.count_float4(program.uops) == (2, 1)
 
@@ -35,8 +36,8 @@ class TestFloat4(unittest.TestCase):
     c = a + b
 
     s = c.schedule()[0]
-    uops = get_program(s.ast, renderer=Device[Device.DEFAULT].renderer,
-                       opts=[Opt(op=OptOps.UPCAST, axis=0, arg=4), Opt(op=OptOps.UPCAST, axis=0, arg=2)]).uops
+    uops = get_program(replace_opts(s.ast, [Opt(op=OptOps.UPCAST, axis=0, arg=4), Opt(op=OptOps.UPCAST, axis=0, arg=2)]),
+                       renderer=Device[Device.DEFAULT].renderer).uops
     assert TestFloat4.count_float4(uops) == (4, 2)
 
   @unittest.skipUnless(Device.DEFAULT in {"CPU"} and AMX, "Only CPU with AMX upcasts float up to size 16")
@@ -47,8 +48,8 @@ class TestFloat4(unittest.TestCase):
       c = a + b
 
       s = c.schedule()[0]
-      return get_program(s.ast, renderer=Device[Device.DEFAULT].renderer,
-                         opts=[Opt(op=OptOps.UPCAST, axis=0, arg=4), Opt(op=OptOps.UPCAST, axis=0, arg=shift)]).uops
+      return get_program(replace_opts(s.ast, [Opt(op=OptOps.UPCAST, axis=0, arg=4), Opt(op=OptOps.UPCAST, axis=0, arg=shift)]),
+                         renderer=Device[Device.DEFAULT].renderer).uops
 
     sizes = [12, 8, 16]
     shifts = [3, 2, 4]
@@ -66,7 +67,7 @@ class TestFloat4(unittest.TestCase):
     s = c.schedule()[0]
     realized_ast = s.ast
     opts_to_apply = [Opt(op=OptOps.UPCAST, axis=0, arg=4)]
-    program = get_program(realized_ast, renderer=Device[Device.DEFAULT].renderer, opts=opts_to_apply)
+    program = get_program(replace_opts(realized_ast, opts_to_apply), renderer=Device[Device.DEFAULT].renderer)
 
     assert TestFloat4.count_float4(program.uops) == (0, 1)
 
@@ -77,8 +78,8 @@ class TestFloat4(unittest.TestCase):
     c = a + b
 
     s = c.schedule()[0]
-    uops = get_program(s.ast, renderer=Device[Device.DEFAULT].renderer,
-                       opts=[Opt(op=OptOps.UPCAST, axis=1, arg=4), Opt(op=OptOps.UPCAST, axis=1, arg=2)]).uops
+    uops = get_program(replace_opts(s.ast, [Opt(op=OptOps.UPCAST, axis=1, arg=4), Opt(op=OptOps.UPCAST, axis=1, arg=2)]),
+                       renderer=Device[Device.DEFAULT].renderer).uops
 
     assert TestFloat4.count_float4(uops) == (0, 2)
 
@@ -90,8 +91,8 @@ class TestFloat4(unittest.TestCase):
       c = a + b
 
       s = c.schedule()[0]
-      return get_program(s.ast, renderer=Device[Device.DEFAULT].renderer,
-                         opts=[Opt(op=OptOps.UPCAST, axis=1, arg=4), Opt(op=OptOps.UPCAST, axis=1, arg=shift)]).uops
+      return get_program(replace_opts(s.ast, [Opt(op=OptOps.UPCAST, axis=1, arg=4), Opt(op=OptOps.UPCAST, axis=1, arg=shift)]),
+                         renderer=Device[Device.DEFAULT].renderer).uops
 
     sizes = [13, 9, 17]
     shifts = [3, 2, 4]
@@ -109,7 +110,7 @@ class TestFloat4(unittest.TestCase):
     # float4 should be emitted (the reduce axis of size 4 is the float4 axis here)
 
     s = c.schedule()[0]
-    uops = get_program(s.ast, renderer=Device[Device.DEFAULT].renderer, opts=[Opt(op=OptOps.UNROLL, axis=0, arg=4)]).uops
+    uops = get_program(replace_opts(s.ast, [Opt(op=OptOps.UNROLL, axis=0, arg=4)]), renderer=Device[Device.DEFAULT].renderer).uops
 
     assert TestFloat4.count_float4(uops) == (0, 0)
 
@@ -123,8 +124,8 @@ class TestFloat4(unittest.TestCase):
     # UPDATE: now we do this fusion
 
     s = c.schedule()[0]
-    uops = get_program(s.ast, renderer=Device[Device.DEFAULT].renderer,
-                       opts=[Opt(op=OptOps.UPCAST, axis=0, arg=0), Opt(op=OptOps.UNROLL, axis=0, arg=0)]).uops
+    uops = get_program(replace_opts(s.ast, [Opt(op=OptOps.UPCAST, axis=0, arg=0), Opt(op=OptOps.UNROLL, axis=0, arg=0)]),
+                       renderer=Device[Device.DEFAULT].renderer).uops
 
     assert TestFloat4.count_float4(uops) in {(0,1), (1,1)}
 
@@ -137,7 +138,7 @@ class TestFloat4(unittest.TestCase):
     # since the top axis is not contiguous.
 
     s = c.schedule()[0]
-    uops = get_program(s.ast, renderer=Device[Device.DEFAULT].renderer, opts=[Opt(op=OptOps.UPCAST, axis=0, arg=4)]).uops
+    uops = get_program(replace_opts(s.ast, [Opt(op=OptOps.UPCAST, axis=0, arg=4)]), renderer=Device[Device.DEFAULT].renderer).uops
 
     assert TestFloat4.count_float4(uops) == (0, 1)
 
@@ -149,7 +150,7 @@ class TestFloat4(unittest.TestCase):
     # should float4 b but not a
 
     s = c.schedule()[0]
-    uops = get_program(s.ast, renderer=Device[Device.DEFAULT].renderer, opts=[Opt(op=OptOps.UPCAST, axis=0, arg=4)]).uops
+    uops = get_program(replace_opts(s.ast, [Opt(op=OptOps.UPCAST, axis=0, arg=4)]), renderer=Device[Device.DEFAULT].renderer).uops
 
     assert TestFloat4.count_float4(uops) == (1, 1)
 

test/opt/test_tensor_cores.py

@@ -8,7 +8,7 @@ from tinygrad.uop.ops import Ops
 from tinygrad.dtype import DType
 from tinygrad.device import is_dtype_supported
 from tinygrad.helpers import AMX, DEV, Context
-from test.helpers import slow
+from test.helpers import slow, replace_opts
 from tinygrad.engine.realize import CompiledRunner, get_program
 from tinygrad.codegen.opt import Opt, OptOps, KernelOptError
 from tinygrad.codegen.opt.tc import amd_cdna_1616128
@@ -27,14 +27,14 @@ def helper_tc_ensure_uops_and_opts_count(N: int, M:int, K:int, dtype_in:DType, d
   opts_to_apply = [Opt(OptOps.TC, axis, (tc_select, tc_opt, 1))]
 
   if ensure_triggered:
-    program = get_program(realized_ast, Device[Device.DEFAULT].renderer, opts=opts_to_apply)
+    program = get_program(replace_opts(realized_ast, opts_to_apply), Device[Device.DEFAULT].renderer)
     wmmas = len([uop for uop in program.uops if uop.op is Ops.WMMA])
     tcs = len([x for x in program.applied_opts if x.op is OptOps.TC])
     assert wmmas > 0, "tensor core not triggered"
     assert tcs == 1, "tensor core opt not included"
   else:
     try:
-      program = get_program(realized_ast, Device[Device.DEFAULT].renderer, opts=opts_to_apply)
+      program = get_program(replace_opts(realized_ast, opts_to_apply), Device[Device.DEFAULT].renderer)
       assert False, "OptOps.TC triggered, expected KernelOptError"
     except KernelOptError: pass
 
@@ -45,7 +45,7 @@ def helper_tc_allclose(N:int, M:int, K:int, dtype_in:DType, dtype_out:DType, axi
   if dtype_in == dtypes.bfloat16: r = r.float()
   realized_ast, bufs = helper_realized_ast(r)
   opts = [Opt(op=OptOps.TC, axis=axis, arg=(tc_select, tc_opt, use_tensor_cores))]
-  prg = CompiledRunner(replace(get_program(realized_ast, Device[Device.DEFAULT].renderer, opts=opts), device=Device.DEFAULT))
+  prg = CompiledRunner(replace(get_program(replace_opts(realized_ast, opts), Device[Device.DEFAULT].renderer), device=Device.DEFAULT))
   if use_tensor_cores == 1: assert len([uop for uop in prg.p.uops if uop.op is Ops.WMMA]) > 0, "wmma not triggered"
   assert len([x for x in prg.p.uops[-1].arg.applied_opts if x.op is OptOps.TC]) == 1, "tensor core opt not included"
   prg.exec(bufs)
@@ -74,7 +74,7 @@ class TestTensorCores(unittest.TestCase):
       n, m, k = tc.dims[0], tc.dims[1], 2 if AMX else tc.dims[2]
       a, b = Tensor.rand(m, k, dtype=tc.dtype_in), Tensor.rand(k, n, dtype=tc.dtype_in)
       r = a.matmul(b, dtype=tc.dtype_out)
-      prg = get_program(r.schedule()[-1].ast, Device[Device.DEFAULT].renderer, opts=[Opt(op=OptOps.TC, axis=0, arg=(-1, 2, 1))])
+      prg = get_program(replace_opts(r.schedule()[-1].ast, [Opt(op=OptOps.TC, axis=0, arg=(-1, 2, 1))]), Device[Device.DEFAULT].renderer)
       if Device.DEFAULT == "CPU" and DEV.renderer == "LLVM":
         assert "0x201000" in prg.src
       elif Device.DEFAULT == "AMD" and DEV.renderer == "LLVM":
@@ -141,7 +141,7 @@ class TestTensorCores(unittest.TestCase):
         c = a.conv2d(b, padding=1, dtype=tc.dtype_out)
         realized_ast, real_bufs = helper_realized_ast(c)
 
-        program = get_program(realized_ast, Device[Device.DEFAULT].renderer, opts=[Opt(OptOps.TC, axis, (-1, 2, 1))])
+        program = get_program(replace_opts(realized_ast, [Opt(OptOps.TC, axis, (-1, 2, 1))]), Device[Device.DEFAULT].renderer)
         assert len([uop for uop in program.uops if uop.op is Ops.WMMA]) > 0, "tensor core not triggered"
         assert len([x for x in program.applied_opts if x.op is OptOps.TC]) == 1, "tensor core opt not included"
 
@@ -165,7 +165,7 @@ class TestTensorCores(unittest.TestCase):
     r = x.matmul(y, dtype=tc.dtype_out)
     opts = [Opt(OptOps.UNROLL, 0, 4)]
     ast = helper_linearizer_opt(r, [opts], apply_tc=True, atol=3e-2, rtol=1e-3)
-    for u in get_program(ast, Device[Device.DEFAULT].renderer, opts=opts).uops:
+    for u in get_program(replace_opts(ast, opts), Device[Device.DEFAULT].renderer).uops:
       if u.op is Ops.WMMA:
         assert u.src[-1].src[0].op != Ops.STORE
 
@@ -179,7 +179,7 @@ class TestTensorCores(unittest.TestCase):
     r = x.matmul(y, dtype=tc.dtype_out)
     opts = [Opt(OptOps.UNROLL, 0, 4)]
     ast = helper_linearizer_opt(r, [opts], apply_tc=True, atol=3e-2, rtol=1e-3)
-    for u in get_program(ast, Device[Device.DEFAULT].renderer, opts=opts).uops:
+    for u in get_program(replace_opts(ast, opts), Device[Device.DEFAULT].renderer).uops:
       if u.op is Ops.WMMA:
         #assert u.src[-1].dtype == dtypes.float.vec(prod(tc.thread_local_sizes[2]))
         assert u.src[-1].src[0].op != Ops.STORE
@@ -195,7 +195,7 @@ class TestTensorCores(unittest.TestCase):
     r = x.matmul(y, dtype=tc.dtype_out).relu()
     opts = [Opt(OptOps.UNROLL, 0, 4)]
     ast = helper_linearizer_opt(r, [opts], apply_tc=True, atol=3e-2, rtol=1e-3)
-    for u in get_program(ast, Device[Device.DEFAULT].renderer, opts=opts).uops:
+    for u in get_program(replace_opts(ast, opts), Device[Device.DEFAULT].renderer).uops:
       if u.op is Ops.WMMA:
         #assert u.src[-1].dtype == dtypes.float.vec(prod(tc.thread_local_sizes[2]))
         assert u.src[-1].src[0].op != Ops.STORE

test/speed/external_test_speed_v_torch.py

@@ -12,7 +12,7 @@ import numpy as np
 np.set_printoptions(linewidth=160)
 from tinygrad import Tensor, Device, GlobalCounters, TinyJit
 from tinygrad.nn import Conv2d
-from tinygrad.helpers import colorize_float, getenv, CI
+from tinygrad.helpers import colorize_float, getenv, CI, DEV
 
 IN_CHANS = [int(x) for x in getenv("IN_CHANS", "4,16,64").split(",")]
 
@@ -113,7 +113,7 @@ def helper_test_conv(bs, in_chans, out_chans, kernel_size, img_size_y, img_size_
   helper_test_generic(f"conv bs:{bs:3d} chans:{in_chans:3d} -> {out_chans:3d} k:{kernel_size}", f1, (torch_dat,), TinyJit(f2), (tiny_dat,))
 
 @unittest.skipIf(getenv("BIG") == 0, "no big tests")
-@unittest.skipIf(getenv("MOCKGPU"), "no MOCKGPUs")
+@unittest.skipIf(DEV.interface.startswith("MOCK"), "no MOCKGPUs")
 class TestBigSpeed(unittest.TestCase):
   def test_add(self):
     def f(a, b): return a+b
@@ -134,7 +134,7 @@ class TestBigSpeed(unittest.TestCase):
   def test_matvec_16384_4096(self): helper_test_matvec('matvec_16384_4096', 16384, 4096)
 
 @unittest.skipIf(getenv("BIG") == 1, "only big tests")
-@unittest.skipIf(getenv("MOCKGPU"), "no MOCKGPUs")
+@unittest.skipIf(DEV.interface.startswith("MOCK"), "no MOCKGPUs")
 class TestSpeed(unittest.TestCase):
   def test_sub(self):
     def f(a, b): return a-b

test/testextra/test_mockgpu.py

@@ -1,7 +1,7 @@
-from tinygrad.helpers import getenv
+from tinygrad.helpers import DEV
 import unittest, importlib
 
-@unittest.skipUnless(getenv("MOCKGPU"), 'Testing mockgpu')
+@unittest.skipUnless(DEV.interface.startswith("MOCK"), 'Testing mockgpu')
 class TestMockGPU(unittest.TestCase):
   # https://github.com/tinygrad/tinygrad/pull/7627
   def test_import_typing_extensions(self):

test/unit/test_assign.py

@@ -130,19 +130,15 @@ class TestAssign(unittest.TestCase):
     new = a + old_a
     np.testing.assert_allclose(new.numpy(), 4)
 
-  @unittest.skip("TODO: this is broken")
   def test_assign_changes_alt(self, realize=False):
     a = Tensor(1).contiguous()
     if realize: a.realize()
-    b = a.contiguous()    # b returns a new Tensor
+    b = a.clone()
     b.assign(2)
     b.realize()
     self.assertNotEqual(a.item(), b.item())
-  # on a realized Tensor contiguous child changes the source
-  @unittest.expectedFailure
   def test_assign_changes_realized_alt(self): return self.test_assign_changes_alt(realize=True)
 
-  @unittest.skip("assign to contiguous shouldn't change the base buffer")
   def test_assign_changes_buffer_alt(self):
     a, b = [Tensor(Tensor(0).contiguous().realize().uop.buf_uop) for _ in range(2)]
     Tensor.realize(a.contiguous().assign(1), b.contiguous().assign(2))
@@ -510,17 +506,6 @@ class TestAssign(unittest.TestCase):
       # TODO: broken now
       np.testing.assert_equal(a.numpy(), [0]*8)
 
-  @unittest.skip("don't use output buffer, and mismatch dtype no longer supported")
-  def test_cast_assignment(self):
-    a = Tensor(np.arange(N*N, dtype=np.float32)).reshape(N,N)
-    a.realize()
-    oba1 = a.uop.base.output_buffer
-    a.assign(a.cast(dtypes.int32).realize())
-    a.realize()
-    oba2 = a.uop.base.output_buffer
-    assert oba1 is None and oba2 is None
-    np.testing.assert_allclose(a.numpy(), np.arange(N*N,dtype=np.int32).reshape((N,N)))
-
   def test_assign_dtype_mismatch(self):
     # assign should not implicitly cast dtypes - this can lose precision
     a = Tensor.zeros(4, dtype=dtypes.float32).contiguous().realize()
@@ -621,6 +606,35 @@ class TestAssign(unittest.TestCase):
     # N matmuls + N assigns + 1 final read = 2*N+1 (AFTER embedding allows full graph scheduling with shared contiguous reuse)
     self.assertEqual(GlobalCounters.kernel_count, 2*N+1)
 
+  def test_double_assign_from_const(self):
+    a = Tensor.empty(2)
+    a.assign(Tensor.ones(2))
+    a.assign(Tensor.ones(2))
+    GlobalCounters.reset()
+    a.realize()
+    self.assertEqual(GlobalCounters.kernel_count, 1)
+    self.assertEqual(a.tolist(), [1.,1.])
+
+  def test_nested_after_contiguous_store(self):
+    # Mirrors the nested contiguous-write-then-assign-back shape from torch backend view updates.
+    base = Tensor.empty(3, dtype=dtypes.int64)
+    base.assign(Tensor([1, 2, 3], dtype=dtypes.int64))
+    contig = base.contiguous()
+    contig.assign(Tensor([1, 4, 3], dtype=dtypes.int64))
+    GlobalCounters.reset()
+    base.assign(contig).realize()
+    self.assertEqual(GlobalCounters.kernel_count, 2)  # TODO: first copy is dead, could be 1
+    self.assertEqual(base.tolist(), [1,4,3])
+
+  def test_nested_after_contiguous_store_no_init(self):
+    # Same shape as test_nested_after_contiguous_store, but without the initial assign.
+    base = Tensor.empty(3, dtype=dtypes.int64)
+    contig = base.contiguous()
+    contig.assign(Tensor([1, 4, 3], dtype=dtypes.int64))
+    GlobalCounters.reset()
+    base.assign(contig).realize()
+    self.assertEqual(GlobalCounters.kernel_count, 1)
+    self.assertEqual(base.tolist(), [1,4,3])
 
 class TestAssignOrdering(unittest.TestCase):
   """Tests for complex assign orderings that could differ between lazy and eager execution.
@@ -658,7 +672,6 @@ class TestAssignOrdering(unittest.TestCase):
     self.assertEqual(r1.item(), 4)
     self.assertEqual(r2.item(), 8)
 
-  @unittest.skip("TODO: this is broken")
   def test_write_read_write_chain(self):
     """Write, read, write chain - middle read must complete before second write."""
     buf = Tensor.zeros(4).contiguous().realize()
@@ -668,7 +681,11 @@ class TestAssignOrdering(unittest.TestCase):
     final_sum = buf.sum()  # lazy read, should be 20
     # Realize in "wrong" order - final first
     self.assertEqual(final_sum.realize().item(), 20)
-    self.assertEqual(mid_sum.realize().item(), 12)
+    try:
+      self.assertEqual(mid_sum.realize().item(), 12)
+    except AssertionError:
+      # TODO: this is wrong
+      self.assertEqual(mid_sum.realize().item(), 20)
 
   def test_slice_read_then_full_write(self):
     """Read from slice, then overwrite full buffer - WAR dependency works for full buffer assigns."""
@@ -958,11 +975,10 @@ class TestAfterCachePatterns(unittest.TestCase):
     a_store = a.uop.store(c.uop)
     b_store = b.uop.store(c.uop)
 
-    with self.assertRaises(AssertionError):
-      a = Tensor(a.uop.after(a_store, b_store))
-      a.realize()
-      np.testing.assert_array_equal(a.numpy(), 1)
-      np.testing.assert_array_equal(b.numpy(), 1)
+    a = Tensor(a.uop.after(a_store, b_store))
+    a.realize()
+    np.testing.assert_array_equal(a.numpy(), 1)
+    np.testing.assert_array_equal(b.numpy(), 1)
 
   def test_double_store_after_different_sizes(self):
     full = Tensor.zeros(2).contiguous()
@@ -974,11 +990,10 @@ class TestAfterCachePatterns(unittest.TestCase):
     full_store = full.uop.store(full_src.uop)
     head_store = head.uop.store(head_src.uop)
 
-    with self.assertRaises(AssertionError):
-      head = Tensor(head.uop.after(head_store, full_store))
-      head.realize()
-      np.testing.assert_array_equal(head.numpy(), [3])
-      np.testing.assert_array_equal(full.numpy(), [1, 2])
+    head = Tensor(head.uop.after(head_store, full_store))
+    head.realize()
+    np.testing.assert_array_equal(head.numpy(), [3])
+    np.testing.assert_array_equal(full.numpy(), [1, 2])
 
 if __name__ == "__main__":
   unittest.main()

test/unit/test_attention.py

@@ -1,7 +1,7 @@
 import unittest
 import numpy as np
 from tinygrad import Tensor, dtypes
-from tinygrad.apps.llm import (
+from tinygrad.llm.model import (
   GatedDeltaNetBlock, SSMConfig, TransformerBlock, TransformerConfig,
   apply_rope as apply_rope_new, precompute_freqs_cis, pairwise_topk,
 )
@@ -70,11 +70,14 @@ class TestGatedDeltaNetBlock(unittest.TestCase):
     return block._attention(x_norm, start_pos).realize().numpy()
 
   def _cache_views(self, block:GatedDeltaNetBlock) -> tuple[np.ndarray, np.ndarray]:
-    conv_flat = (block.ssm_conv_kernel - 1) * block.conv_channels
-    cache = block.delta_cache.numpy()
-    conv_state = cache[:, :conv_flat].reshape(cache.shape[0], block.ssm_conv_kernel - 1, block.conv_channels)
-    recurrent_state = cache[:, conv_flat:].reshape(cache.shape[0], block.num_v_heads, block.head_v_dim, block.head_v_dim)
-    return conv_state, recurrent_state
+    if hasattr(block, 'conv_state'):
+      return block.conv_state.numpy(), block.recurrent_state.numpy()
+    else:
+      conv_flat = (block.ssm_conv_kernel - 1) * block.conv_channels
+      cache = block.delta_cache.numpy()
+      conv_state = cache[:, :conv_flat].reshape(cache.shape[0], block.ssm_conv_kernel - 1, block.conv_channels)
+      recurrent_state = cache[:, conv_flat:].reshape(cache.shape[0], block.num_v_heads, block.head_v_dim, block.head_v_dim)
+      return conv_state, recurrent_state
 
   def _linear_np(self, x:np.ndarray, weight:np.ndarray) -> np.ndarray:
     return x.astype(np.float32) @ weight.T.astype(np.float32)

test/unit/test_call.py

@@ -60,7 +60,6 @@ class TestCall(unittest.TestCase):
     c = Tensor.call(a, b, fxn=a.as_param(0) @ b.as_param(1))
     np.testing.assert_allclose(c.numpy(), a.numpy() @ b.numpy(), rtol=1e-5, atol=1e-6)
 
-  @unittest.skip("needs GEMM on mixins")
   def test_call_gemm_uop(self):
     M, K, N = 4, 8, 4
     a = Tensor.randn(M, K)
@@ -220,9 +219,9 @@ class TestCallSchedule(unittest.TestCase):
     a = Tensor.empty(4, 8)
     b = Tensor.empty(4, 8)
     r0, r1 = f(a), f(b)
-    # find the CALL nodes
-    c0 = next(u for u in r0.uop.toposort() if u.op is Ops.CALL)
-    c1 = next(u for u in r1.uop.toposort() if u.op is Ops.CALL)
+    # find the FUNCTION nodes
+    c0 = next(u for u in r0.uop.toposort() if u.op is Ops.FUNCTION)
+    c1 = next(u for u in r1.uop.toposort() if u.op is Ops.FUNCTION)
     # the function bodies (src[0]) should have identical keys — unique consts must not leak through
     self.assertEqual(c0.src[0].key, c1.src[0].key)
 

test/unit/test_cpu.py

@@ -0,0 +1,24 @@
+import unittest, io
+from contextlib import redirect_stdout
+from tinygrad import Tensor, Device
+from tinygrad.helpers import Target
+from tinygrad.renderer.nir import LVPRenderer
+from tinygrad.engine.realize import get_program
+
+@unittest.skipIf(Device.DEFAULT != "CPU", "only run on CPU")
+class TestCPU(unittest.TestCase):
+  def test_arch_feats(self):
+    ast = (Tensor.empty(16) + Tensor.empty(16)).schedule()[-1].ast
+    for ren in Device[Device.DEFAULT].renderers:
+      for arch, expect_vmov in [("x86_64,x86-64,avx", True), ("x86_64,x86-64,-avx", False)]:
+        with self.subTest(arch=arch):
+          if ren is LVPRenderer: continue # LVP does not play nice with cross compilation
+          r = ren(Target(device="CPU", arch=arch))
+          p = get_program(ast, r)
+          lib = r.compiler.compile(p.src)
+          out = io.StringIO()
+          with redirect_stdout(out): r.compiler.disassemble(lib)
+          self.assertEqual("vmov" in out.getvalue(), expect_vmov, out.getvalue())
+
+if __name__ == '__main__':
+  unittest.main()

test/unit/test_gguf.py

@@ -1,12 +1,31 @@
-import os, struct, unittest
+import os, struct, unittest, sys
 from tinygrad import dtypes, Tensor, fetch, Device
-from tinygrad.nn.state import ggml_data_to_tensor, gguf_load
+from tinygrad.helpers import disable_gc
+from tinygrad.llm.gguf import _ggml_iq_grid, ggml_data_to_tensor, gguf_load
+from tinygrad.runtime.autogen import ggml_common as _ggml
 from tinygrad.device import is_dtype_supported
 import numpy as np
 from gguf import GGUFReader, GGUFValueType, GGMLQuantizationType, GGML_QUANT_SIZES, dequantize, quantize
+from gguf.quants import IQ2_S, IQ3_S, IQ3_XXS
 
 ggml_test_block_count = 4
 
+class TestGGUFTables(unittest.TestCase):
+  def test_iq2_s_grid_matches_gguf_py(self):
+    IQ2_S.init_grid()
+    grid = _ggml_iq_grid(Device.DEFAULT, _ggml.iq2s_grid, (1024, 8)).numpy()
+    np.testing.assert_equal(grid, IQ2_S.grid.reshape(1024, 8))
+
+  def test_iq3_xxs_grid_matches_gguf_py(self):
+    IQ3_XXS.init_grid()
+    grid = _ggml_iq_grid(Device.DEFAULT, _ggml.iq3xxs_grid, (256, 4)).numpy()
+    np.testing.assert_equal(grid, IQ3_XXS.grid.reshape(256, 4))
+
+  def test_iq3_s_grid_matches_gguf_py(self):
+    IQ3_S.init_grid()
+    grid = _ggml_iq_grid(Device.DEFAULT, _ggml.iq3s_grid, (512, 4)).numpy()
+    np.testing.assert_equal(grid, IQ3_S.grid.reshape(512, 4))
+
 @unittest.skipIf(any(not is_dtype_supported(t) for t in [ dtypes.uint8, dtypes.half ]), "Backend must support uint8 and half")
 class TestGGUF(unittest.TestCase):
   def test_load_tinyllama_q8_0(self): self._test_gguf_load("https://huggingface.co/ggml-org/models/resolve/main/tinyllamas/stories15M-q8_0.gguf?download=true")
@@ -36,6 +55,10 @@ class TestGGUF(unittest.TestCase):
   def test_dequantization_q4_k(self): self._test_dequantization(GGMLQuantizationType.Q4_K)
   def test_dequantization_q5_k(self): self._test_dequantization(GGMLQuantizationType.Q5_K)
   def test_dequantization_q6_k(self): self._test_dequantization(GGMLQuantizationType.Q6_K)
+  def test_dequantization_iq3_xxs(self): self._test_dequantization(GGMLQuantizationType.IQ3_XXS)
+  def test_dequantization_iq3_s(self): self._test_dequantization(GGMLQuantizationType.IQ3_S)
+  def test_dequantization_iq2_s(self): self._test_dequantization(GGMLQuantizationType.IQ2_S)
+  def test_dequantization_iq4_xs(self): self._test_dequantization(GGMLQuantizationType.IQ4_XS)
   def test_dequantization_mxfp4(self): self._test_dequantization(GGMLQuantizationType.MXFP4)
   @unittest.skipUnless(is_dtype_supported(dtypes.bfloat16), "Backend must support bfloat16")
   def test_dequantization_bf16(self): self._test_dequantization(GGMLQuantizationType.BF16)
@@ -140,7 +163,10 @@ class TestGGUFGEMV(unittest.TestCase):
       # generate random quantized blocks with valid fp16 scale fields (random bytes can produce NaN scales)
       q_data = rng.integers(0, 256, size=n_blocks * type_size, dtype=np.uint8).reshape(n_blocks, type_size)
       scales = np.float16(rng.standard_normal(n_blocks * 4)).view(np.uint8).reshape(n_blocks, -1)
-      if qtype in (GGMLQuantizationType.Q5_0, GGMLQuantizationType.Q8_0): q_data[:, :2] = scales[:, :2]  # d at offset 0
+      if qtype in (GGMLQuantizationType.Q5_0, GGMLQuantizationType.Q8_0,
+                   GGMLQuantizationType.IQ3_XXS,
+                   GGMLQuantizationType.IQ2_S,
+                   GGMLQuantizationType.IQ3_S, GGMLQuantizationType.IQ4_XS): q_data[:, :2] = scales[:, :2]  # d at offset 0
       elif qtype in (GGMLQuantizationType.Q5_1, GGMLQuantizationType.Q4_K, GGMLQuantizationType.Q5_K):
         q_data[:, :4] = scales[:, :4]  # d, m/dmin at offset 0
       elif qtype == GGMLQuantizationType.Q6_K: q_data[:, -2:] = scales[:, :2]               # d at end
@@ -173,9 +199,23 @@ class TestGGUFGEMV(unittest.TestCase):
   def test_gguf_gemv_q4_k(self): self._test_gguf_gemv(GGMLQuantizationType.Q4_K)
   def test_gguf_gemv_q5_k(self): self._test_gguf_gemv(GGMLQuantizationType.Q5_K)
   def test_gguf_gemv_q6_k(self): self._test_gguf_gemv(GGMLQuantizationType.Q6_K)
+  def test_gguf_gemv_iq3_xxs(self): self._test_gguf_gemv(GGMLQuantizationType.IQ3_XXS)
+  def test_gguf_gemv_iq3_s(self): self._test_gguf_gemv(GGMLQuantizationType.IQ3_S)
+  def test_gguf_gemv_iq2_s(self): self._test_gguf_gemv(GGMLQuantizationType.IQ2_S)
+  def test_gguf_gemv_iq4_xs(self): self._test_gguf_gemv(GGMLQuantizationType.IQ4_XS)
   def test_gguf_gemv_mxfp4(self): self._test_gguf_gemv(GGMLQuantizationType.MXFP4)
   @unittest.skipUnless(is_dtype_supported(dtypes.bfloat16), "Backend must support bfloat16")
   def test_gguf_gemv_bf16(self): self._test_gguf_gemv(GGMLQuantizationType.BF16)
 
+class TestGGUFGC(unittest.TestCase):
+  def test_gguf_load_no_tensor_leak(self):
+    """gguf_load must not retain references to the input tensor after returning."""
+    fp = fetch("https://huggingface.co/ggml-org/models/resolve/main/tinyllamas/stories15M-q8_0.gguf?download=true")
+    t = Tensor.empty(os.stat(fp).st_size, dtype=dtypes.uint8, device=f"disk:{fp}").to(Device.DEFAULT).realize()
+    with disable_gc():
+      ref_before = sys.getrefcount(t)
+      kv_data, tensors = gguf_load(t)
+      self.assertEqual(sys.getrefcount(t), ref_before, "gguf_load leaked a reference to the input tensor")
+
 if __name__ == '__main__':
   unittest.main()

test/unit/test_hashing.py

@@ -1,13 +1,14 @@
 from typing_extensions import Callable
 import hashlib, random, unittest
-from tinygrad import Tensor, Device, getenv, dtypes
+from tinygrad import Tensor, Device, dtypes
+from tinygrad.helpers import DEV
 from test.helpers import slow
 from tinygrad.device import is_dtype_supported
 from tinygrad.uop.ops import UOp
 from tinygrad.engine.jit import TinyJit
 
 @unittest.skipUnless(is_dtype_supported(dtypes.uint8) and is_dtype_supported(dtypes.uint64), "Device must support uint8 and uint64")
-@unittest.skipIf(getenv("MOCKGPU") and Device.DEFAULT == "NV", "crashes in NV CI")
+@unittest.skipIf(DEV.interface.startswith("MOCK") and Device.DEFAULT == "NV", "crashes in NV CI")
 class TestHashing(unittest.TestCase):
   def _python_hash_1mb(self, data:bytes):
     chunks = [data[i:i+4096] for i in range(0, len(data), 4096)]
@@ -21,7 +22,7 @@ class TestHashing(unittest.TestCase):
     self.assertEqual(bytes(out.data()), expected)
 
 @unittest.skipUnless(is_dtype_supported(dtypes.uint8) and is_dtype_supported(dtypes.uint64), "Device must support uint8 and uint64")
-@unittest.skipIf(getenv("MOCKGPU") and Device.DEFAULT == "NV", "crashes in NV CI")
+@unittest.skipIf(DEV.interface.startswith("MOCK") and Device.DEFAULT == "NV", "crashes in NV CI")
 class TestKeccak(unittest.TestCase):
   def setUp(self) -> None: random.seed(1337)
 

test/unit/test_llm_mla.py

@@ -1,7 +1,7 @@
 import unittest
 import numpy as np
 from tinygrad import Tensor
-from tinygrad.apps.llm import Transformer, TransformerConfig, apply_rope
+from tinygrad.llm.model import Transformer, TransformerConfig, apply_rope, MLATransformerBlock, precompute_freqs_cis
 
 class TestMLA(unittest.TestCase):
   def _make_config(self, **kwargs):
@@ -13,7 +13,6 @@ class TestMLA(unittest.TestCase):
 
   def test_mla_attention_matches_naive(self):
     config = self._make_config(max_context=16)
-    from tinygrad.apps.llm import MLATransformerBlock, precompute_freqs_cis
 
     block = MLATransformerBlock(config)
     c = config

test/unit/test_llm_moe.py

@@ -2,7 +2,7 @@ import unittest
 import numpy as np
 from dataclasses import replace
 from tinygrad import Tensor
-from tinygrad.apps.llm import TransformerBlock, TransformerConfig
+from tinygrad.llm.model import TransformerBlock, TransformerConfig
 
 def _moe_config(dim=8, hidden=16, n_heads=2, num_experts=4, num_experts_per_tok=2):
   return TransformerConfig(

test/unit/test_llm_server.py

@@ -2,7 +2,7 @@ import unittest
 from unittest.mock import patch
 from tinygrad import Tensor, UOp
 from tinygrad.schedule import schedule_cache
-from tinygrad.apps.llm import Transformer, TransformerConfig
+from tinygrad.llm.model import Transformer, TransformerConfig
 
 TEST_CONFIG = TransformerConfig(num_blocks=1, dim=64, hidden_dim=128, n_heads=2, n_kv_heads=2,
                            norm_eps=1e-5, vocab_size=100, head_dim=32, rope_theta=10000.0, rope_dim=32, v_head_dim=32, max_context=32)

test/unit/test_objc.py

@@ -0,0 +1,17 @@
+import sys, unittest
+
+class TestObjCMetaSpec(unittest.TestCase):
+  @unittest.skipUnless(sys.platform == "darwin", "objc runtime only on macOS")
+  def test_classmethods_are_classmethods(self):
+    from tinygrad.runtime.support.objc import Spec, id_
+
+    #_classmethods_ must include classmethod descriptors
+    class ObjCTest(Spec):
+      _methods_ = [("foo", id_, [])]
+      _classmethods_ = [("bar", id_, [])]
+
+    self.assertNotIsInstance(ObjCTest.__dict__["foo"], classmethod)
+    self.assertIsInstance(ObjCTest.__dict__["bar"], classmethod)
+
+if __name__ == "__main__":
+  unittest.main()

tinygrad/callify.py

@@ -18,7 +18,7 @@ def tag_uop(ctx:AllocCtx, x:UOp):
 def disk_copy_is_buffer(ctx:AllocCtx, u:UOp):
   # copies to disk are replaced with the disk buffer
   to_disk = isinstance(u._device, str) and u._device.startswith(("DISK", "TINYFS"))
-  if to_disk: ctx.buffer_map[u] = UOp.new_buffer(u.device, u.shard_size, u.dtype).reshape(u.max_shard_shape)
+  if to_disk: ctx.buffer_map[u] = u.empty_like()
   # all copies from disk/numpy are realized into a real buffer
   from_creation = isinstance(u.src[0]._device, str) and any(u.src[0]._device.startswith(x) for x in ["NPY", "DISK", "PYTHON", "TINYFS"])
   if from_creation: return tag_uop(ctx, u)
@@ -40,19 +40,14 @@ add_tags = PatternMatcher([
   (UPat(GroupOp.All, name="x"), lambda ctx,x: tag_uop(ctx,x) if x in ctx.bases else None),
 ])
 
-def _buffer_like(u:UOp) -> UOp:
-  buffer = UOp.new_buffer(u.device, u.shard_size, u.dtype).reshape(u.max_shard_shape).shrink_to(u.shard_shape)
-  if isinstance(u.device, tuple) and u.axis is not None: buffer = buffer.multi(u.axis)
-  return buffer
-
 def replace_contig_with_store_after(u:UOp):
   # can't allocate a buffer without a device (e.g., inside a CALL function body with only PARAMs)
   if u._device is None: return None
   # if size is 0, remove the contig
-  if u.size == 0: return u.src[0]
+  if 0 in u.shape: return u.src[0]
   # no real contig for DISK/TINYFS tensors, they are left alone
   if isinstance(u._device, str) and u._device.startswith(("DISK", "TINYFS")): return u.rtag(None)
-  buf = _buffer_like(u)
+  buf = u.empty_like()
   return buf.after(buf.store(u.src[0])).rtag(u.tag)
 
 def replace_store_after_with_contig(u:UOp, src:UOp):
@@ -96,31 +91,30 @@ def contiguous_mops_to_view(c:UOp, src:UOp):
 
 def transform_precompiled_call(c:UOp) -> UOp|None:
   if not c.arg.precompile: return None
-  if c.src[0].op is Ops.SINK: return None
-  assert c.src[0].op is Ops.TUPLE, f"expected TUPLE body for precompiled call, got {c.src[0].op}"
+  assert c.src[0].op is Ops.TUPLE, f"expected TUPLE body for precompiled FUNCTION, got {c.src[0].op}"
   input_buffers = tuple(x.contiguous() if x.op not in {Ops.AFTER, Ops.BIND} else x for x in c.src[1:])
 
   # add the outputs to the call
   srcs = c.src[0].src
   resolved = [c.gettuple(i) for i in range(len(srcs))]
-  outs = tuple(_buffer_like(r) for r in resolved)
+  outs = tuple(r.empty_like() for r in resolved)
   targets = [o.param_like(len(c.src)-1+i).shrink_to(s.shape) for i,(o,s) in enumerate(zip(outs, srcs))]
   fxn = UOp.sink(*[t.after(t.store(s)) for t,s in zip(targets, srcs)])
 
-  # create the new thing for the big graph
-  new_call = c.replace(src=(fxn, *input_buffers, *outs), tag=None)
+  # body switches from TUPLE to SINK, so the node becomes an opaque CALL (not FUNCTION)
+  new_call = UOp(Ops.CALL, c.dtype, (fxn, *input_buffers, *outs), c.arg)
   rets = tuple(o.after(new_call) for o in outs)
 
   # if the CALL has symbolic shapes, shrink the max-sized output to the actual symbolic shape
-  # NOTE: must use resolved shapes from the CALL (which substitutes PARAMs with external args), not raw body shapes
+  # NOTE: must use resolved shapes from the FUNCTION (which substitutes PARAMs with external args), not raw body shapes
   rets = tuple(r.shrink_to(rs.shape) for r,rs in zip(rets, resolved))
 
   return UOp.maketuple(*rets)
 
 # NOTE: adding rules to here is bad. these all need to run before the schedule cache
 pm_early_transform_tensor_graph = PatternMatcher([
-  # transform precompiled CALLs
-  (UPat(Ops.CALL, name="c"), transform_precompiled_call),
+  # transform precompiled FUNCTIONs into CALLs (body becomes SINK with stores)
+  (UPat(Ops.FUNCTION, name="c"), transform_precompiled_call),
 
   # resolve TUPLE+GETTUPLE (for precompiled calls)
   (UPat(Ops.GETTUPLE, src=(UPat(Ops.TUPLE, name="t"),), name="g"), lambda g,t: t.src[g.arg]),
@@ -142,29 +136,28 @@ pm_early_transform_tensor_graph = PatternMatcher([
   (UPat((Ops.DETACH, Ops.CONTIGUOUS_BACKWARD), name="x"), lambda x: x.src[0]),
 ])
 
-def untag_and_append(ctx:AllocCtx, x:UOp):
-  if x.tag is None: return None
+def finalize_after(ctx:AllocCtx, x:UOp):
+  # untagged: record as an assign for the call body
+  if x.tag is None:
+    ctx.assigns.append(x)
+    return None
+  # tagged: untag and map each original pre-rewrite UOp to the stripped buffer; the untagged result is reprocessed as untagged
   ret = x.replace(tag=None)
+  replace_uop = ret
+  while replace_uop.op is Ops.AFTER: replace_uop = replace_uop.src[0]
   for t in x.tag:
     original_uop: UOp = ctx.uop_list[t]
-    replace_uop = ret
-    while replace_uop.op is Ops.AFTER: replace_uop = replace_uop.src[0]
     ctx.buffer_map[original_uop] = replace_uop.shrink_to(original_uop.shape)
-  if ret.op is not Ops.AFTER: ctx.assigns.append(ret)  # AFTER gets appended by append_after
   return ret
 
-def append_after(ctx:AllocCtx, x:UOp):
-  ctx.assigns.append(x)
-
 def replace_input_buffer(ctx:AllocCtx, b:UOp):
   ctx.replacements.append(b)
   return UOp.param(len(ctx.replacements)-1, b.dtype, b.shape, b._device,
                    b._min_max if b.op is Ops.BIND else None, b.src[0].arg[0] if b.op is Ops.BIND else None)
 
 pm_finalize_call = PatternMatcher([
-  (UPat(Ops.AFTER, name="x"), untag_and_append),
-  (UPat(Ops.AFTER, name="x"), append_after),
-  (UPat(Ops.COPY, name="x"), lambda ctx,x: append_after(ctx,x) if isinstance(x.device, str) and x.device.startswith(("DISK", "TINYFS")) else None),
+  (UPat(Ops.AFTER, name="x"), finalize_after),
+  (UPat(Ops.COPY, name="x"), lambda ctx,x: ctx.assigns.append(x) if isinstance(x.device, str) and x.device.startswith(("DISK", "TINYFS")) else None),
   # remove unique from const. TODO: this is copied in function.py
   (UPat(Ops.CONST, src=(UPat(Ops.UNIQUE), UPat(Ops.DEVICE, name="d")), name="b"), lambda b,d: b.replace(src=(d,))),
 ])

tinygrad/codegen/__init__.py

@@ -6,7 +6,6 @@ from tinygrad.uop.ops import PatternMatcher, graph_rewrite, UOp, pm_lower_index_
 from tinygrad.uop.spec import type_verify, program_spec, kernel_spec
 from tinygrad.renderer import Renderer, ProgramSpec, Estimates
 from tinygrad.dtype import dtypes
-from tinygrad.codegen.opt import Opt
 
 # import all pattern matchers here
 from tinygrad.codegen.gpudims import pm_add_gpudims
@@ -151,7 +150,7 @@ pm_to_program = PatternMatcher([
 
 @Context(ALLOW_DEVICE_USAGE=0)
 @track_rewrites(name=lambda ast,renderer,ret,**kwargs: TracingKey(ret.name, (ret.function_name, ast), ret=renderer), replay=True)
-def get_program(ast:UOp, renderer:Renderer, opts:list[Opt]|None=None) -> ProgramSpec:
+def get_program(ast:UOp, renderer:Renderer) -> ProgramSpec:
   """
   Transform an AST into a ProgramSpec. May trigger BEAM search.
 
@@ -168,10 +167,6 @@ def get_program(ast:UOp, renderer:Renderer, opts:list[Opt]|None=None) -> Program
     beam, ast = (ast.arg, ast.src[0]) if ast.op is Ops.BEAM else (0, ast)
     # rewrite to prg
     assert isinstance(ast.arg, KernelInfo), "requires KernelInfo on arg to get_program"
-    if opts is not None:
-      # TODO: should this be here?
-      assert ast.arg.opts_to_apply is None, "can't apply opts if there's already opts to apply"
-      ast = ast.replace(arg=replace(ast.arg, opts_to_apply=tuple(opts)))
     full_sink = full_rewrite_to_sink(ast, renderer, optimize=ast.tag is None, beam=beam)
     prg = UOp(Ops.PROGRAM, src=(full_sink, UOp(Ops.DEVICE, arg=renderer.target.device)))
   else:

tinygrad/codegen/gpudims.py

@@ -1,5 +1,5 @@
 import math
-from tinygrad.uop.ops import UOp, Ops, sint, PatternMatcher, UPat, KernelInfo, ssimplify, AxisType, sint_to_uop
+from tinygrad.uop.ops import UOp, Ops, sint, PatternMatcher, UPat, KernelInfo, ssimplify, AxisType
 from tinygrad.helpers import dedup, get_contraction
 from tinygrad.dtype import dtypes, AddrSpace, Invalid
 from tinygrad.renderer import Renderer
@@ -35,7 +35,7 @@ def get_grouped_dims(prefix, dims:tuple[sint, ...], max_sizes:tuple[int, ...]|No
     if len(limited) > len(max_sizes): raise RuntimeError(f"cannot limit dim {dims=}, {max_sizes=}")
     # try to split up dims: (a,) -> (b, c)
     if limited == dims: limited = _split_dims(dims, max_sizes)
-  raw_idxs = [UOp(Ops.SPECIAL, dtypes.weakint, (sint_to_uop(s),), (f"{prefix}{i}")) for i,s in enumerate(limited)]
+  raw_idxs = [UOp.special(s, f"{prefix}{i}") for i,s in enumerate(limited)]
   if len(limited) < len(dims):
     ret = []
     if (contraction:=get_contraction(dims, limited)) is None: raise RuntimeError(f"get_contraction should not be None {dims=} {limited=}")

tinygrad/codegen/late/devectorizer.py

@@ -318,7 +318,7 @@ def reduce_to_acc(ctx:ReduceContext, red:UOp):
     ended_ranges = flatten([x.ended_ranges for x in topo if x.op is Ops.END])
     input_ranges = tuple([x for x in topo if x.op is Ops.RANGE and x not in reduce_range and x not in ended_ranges])
     identity = red.const(red.dtype, identity_element(red.arg, red.dtype.scalar()))
-    acc = UOp(Ops.DEFINE_REG, red.dtype.ptr(size=1, addrspace=AddrSpace.REG), arg=ctx.acc_num)
+    acc = UOp.placeholder((1,), red.dtype, ctx.acc_num, AddrSpace.REG)
     acc_init = acc.after(*input_ranges).index(UOp.const(dtypes.int, 0)).store(identity)
     lst = [acc.after(acc_init, *reduce_range).index(UOp.const(dtypes.int, 0))] + lst  # put acc as the first element
     ctx.acc_num += 1

tinygrad/codegen/opt/postrange.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 import math, itertools
 from collections import defaultdict
 from typing import cast, Final
-from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, KernelInfo, graph_rewrite, AxisType, ssimplify, GroupOp
+from tinygrad.uop.ops import Ops, UOp, KernelInfo, graph_rewrite, AxisType, ssimplify, GroupOp, remove_all_tags
 from tinygrad.uop.ops import axis_letters, axis_colors, axis_to_pos
 from tinygrad.device import Buffer
 from tinygrad.dtype import dtypes
@@ -12,8 +12,6 @@ from tinygrad.codegen.opt import Opt, OptOps, KernelOptError, check
 from tinygrad.codegen.simplify import pm_flatten_range
 from tinygrad.renderer import Renderer
 
-remove_tags = PatternMatcher([(UPat(GroupOp.All, name="x"), lambda x: x.replace(tag=None) if x.tag is not None else None)])
-
 class Scheduler:
   def __init__(self, ast:UOp, ren:Renderer):
     self.ast, self.ren = ast, ren
@@ -211,7 +209,7 @@ class Scheduler:
       self.ast = self.ast.substitute({rng:rng.replace(arg=(*altrng.arg[0:-1], rng.arg[-1]), tag=1),
                                       altrng:altrng.replace(arg=(*rng.arg[0:-1], altrng.arg[-1]), tag=1)},
                                       name=f"swap {rng.arg[:-1]} {altrng.arg[:-1]}")
-      self.ast = graph_rewrite(self.ast, remove_tags, name="swap remove tags")
+      self.ast = graph_rewrite(self.ast, remove_all_tags, name="swap remove tags")
     else:
       raise KernelOptError(f"unsupported opt {opt.op}")
 

tinygrad/codegen/simplify.py

@@ -134,7 +134,7 @@ def reduce_collapse(red:UOp, u:UOp, pm:PatternMatcher=pm_reduce_collapse) -> UOp
     for u in included:
       for s in u.src:
         if s in included or s in replaces or s.op in {Ops.CONST, Ops.VCONST, Ops.PARAM, Ops.DEFINE_LOCAL, Ops.DEFINE_VAR}: continue
-        replaces[s] = UOp(Ops.DEFINE_VAR, dtype=s.dtype, arg=(f'in{len(replaces)}', s.vmin, s.vmax))
+        replaces[s] = UOp.variable(f'in{len(replaces)}', s.vmin, s.vmax, s.dtype)
     collapse_fxn = u.substitute(replaces).reduce(r, arg=Ops.ADD)
     sink = graph_rewrite(collapse_fxn, pm, name="reduce_collapse")
     if not no_range(sink): return None

tinygrad/device.py

@@ -5,7 +5,7 @@ from typing import Any, Generic, TypeVar, Iterator, Generator, TYPE_CHECKING
 import importlib, inspect, functools, pathlib, os, platform, contextlib, sys, re, atexit, pickle, decimal
 from tinygrad.helpers import BENCHMARKS, CI, OSX, LRU, getenv, diskcache_get, diskcache_put, DEBUG, GlobalCounters, flat_mv, PROFILE, temp, colored
 from tinygrad.helpers import Context, CCACHE, ALLOW_DEVICE_USAGE, MAX_BUFFER_SIZE, cpu_events, ProfileEvent, ProfilePointEvent, suppress_finalizing
-from tinygrad.helpers import select_by_name, select_first_inited, DEV, EMULATED_DTYPES, IMAGE, FLOAT16, TracingKey, size_to_str, Target
+from tinygrad.helpers import select_by_name, select_first_inited, DEV, EMULATED_DTYPES, IMAGE, FLOAT16, TracingKey, size_to_str, Target, VIZ
 from tinygrad.dtype import DType, PtrDType, dtypes, _to_np_dtype
 if TYPE_CHECKING: from tinygrad.renderer import Renderer
 
@@ -364,7 +364,7 @@ if PROFILE:
     with open(fn:=temp("profile.pkl", append_user=True), "wb") as f: pickle.dump(cpu_events+Compiled.profile_events+Buffer.profile_events, f)
 
     PROFILE.value = 0
-    if getenv("VIZ") > 0:
+    if VIZ > 0:
       from tinygrad.uop.ops import launch_viz
       launch_viz("PROFILE", fn)
 

tinygrad/engine/jit.py

@@ -1,11 +1,11 @@
-from typing import TypeVar, Generic, Callable, cast, Any
+from typing import TypeVar, Generic, Callable, Any
 import functools, collections
 from tinygrad.tensor import Tensor
-from tinygrad.helpers import flatten, merge_dicts, DEBUG, Context, BEAM, getenv, colored, JIT, JIT_BATCH_SIZE, dedup, unwrap, pluralize, VIZ
+from tinygrad.helpers import flatten, merge_dicts, DEBUG, Context, BEAM, getenv, colored, JIT, JIT_BATCH_SIZE, dedup, pluralize, VIZ
 from tinygrad.device import Buffer, Compiled, Device, MultiBuffer
 from tinygrad.dtype import DType, dtypes
 from tinygrad.uop.ops import UOp, PatternMatcher, Variable, sym_infer, Ops, buffers, track_rewrites, graph_rewrite
-from tinygrad.engine.realize import ExecItem, capturing, BufferCopy, BufferXfer, EncDec, CompiledRunner, Runner, Estimates
+from tinygrad.engine.realize import ExecItem, capturing, CompiledRunner, Runner, Estimates, pm_beam, run_linear, get_runner, graph_cache
 from tinygrad.schedule.memory import memory_plan_rewrite, _collect_bufs
 from tinygrad.schedule import linear_to_schedule
 from tinygrad.nn.state import get_parameters
@@ -22,14 +22,13 @@ def prune_linear(linear:UOp, needed:set[UOp]) -> tuple[UOp, UOp]:
     else: onetime.append(si)
   return linear.replace(src=tuple(kept)), linear.replace(src=tuple(onetime))
 
-def create_graph_call(batch:list[UOp], input_buffers:set[Buffer]) -> UOp:
-  def bufs_for(b): return b.buffer.bufs if isinstance(b.buffer, MultiBuffer) else [b.buffer]
-
-  input_list = dedup(b for si in batch for b in si.src[1:] if b.op in (Ops.BUFFER, Ops.BUFFER_VIEW) and not input_buffers.isdisjoint(bufs_for(b)))
+def create_graph_call(batch:list[UOp]) -> UOp:
+  # all external inputs are PARAMs
+  input_list = dedup(u for si in batch for b in si.src[1:] for u in b.toposort() if u.op is Ops.PARAM)
   cf = UOp(Ops.CUSTOM_FUNCTION, dtypes.void, src=(UOp(Ops.LINEAR, src=tuple(batch)), *input_list), arg="graph")
   return cf.call(*input_list, metadata=tuple(m for si in batch for m in si.arg.metadata))
 
-def graph_split_rewrite(linear:UOp, input_buffers:set[Buffer], max_batch_size:int=0) -> UOp:
+def graph_split_rewrite(linear:UOp, max_batch_size:int=0) -> UOp:
   new_src: list[UOp] = []
   current_batch: list[UOp] = []
   current_batch_devs: list[Compiled] = []
@@ -38,7 +37,7 @@ def graph_split_rewrite(linear:UOp, input_buffers:set[Buffer], max_batch_size:in
     nonlocal current_batch, current_batch_devs, max_batch_size, new_src
     if len(current_batch) <= 1 and not getenv("GRAPH_ONE_KERNEL"): new_src.extend(current_batch)
     else:
-      new_src.append(create_graph_call(current_batch, input_buffers))
+      new_src.append(create_graph_call(current_batch))
       max_batch_size *= 2
       if DEBUG >= 2: print(f"JIT GRAPHing batch with {len(current_batch)} kernels")
     current_batch, current_batch_devs = [], []
@@ -60,19 +59,36 @@ def graph_split_rewrite(linear:UOp, input_buffers:set[Buffer], max_batch_size:in
   if current_batch: flush_batch()
   return linear.replace(src=tuple(new_src))
 
-def jit_cache_bufs(jit_cache:list[ExecItem]):
-  for ei in jit_cache:
-    for b in ei.bufs:
-      if b is not None: yield b
-    if isinstance(ei.prg, GraphRunner): yield from jit_cache_bufs(ei.prg.jit_cache)
+def _unwrap_beam(ast:UOp) -> UOp: return ast.src[0] if ast.op is Ops.BEAM else ast
 
-@track_rewrites(lambda linear,held_bufs,input_buffers=None,ret=(): f"JIT {pluralize('call', len(linear.src))}")
-def jit_lower(linear:UOp, held_bufs:set[UOp], input_buffers:list[Buffer]|None=None) -> list[ExecItem]:
+def _call_outs_ins(call:UOp) -> tuple[set[int], set[int]]:
+  non_bind = [s for s in call.src[1:] if s.op is not Ops.BIND]
+  ast = _unwrap_beam(call.src[0])
+  if ast.op in (Ops.SINK, Ops.PROGRAM):
+    prg = get_runner(non_bind[0].device if isinstance(non_bind[0].device, str) else non_bind[0].device[0], call.src[0])
+    return set(prg.p.outs), set(prg.p.ins)
+  if ast.op in (Ops.COPY, Ops.BUFFER_VIEW): return {0}, {1}
+  if ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "encdec": return {0}, set(range(1, len(non_bind)))
+  return set(), set()
+
+def _copy_input(u:UOp) -> UOp:
+  run_linear(UOp(Ops.LINEAR, src=(u.copy_to_device(u.device).call(new:=UOp.new_buffer(u.device, u.arg, u.dtype), u, metadata=()),)))
+  return new
+
+@track_rewrites(lambda linear,held_bufs,input_uops,ret=(): f"JIT {pluralize('call', len(linear.src))}")
+def jit_lower(linear:UOp, held_bufs:set[UOp], input_uops:list[UOp]) -> UOp:
   if VIZ: graph_rewrite(linear, PatternMatcher([]), name="View captured linear")
+
+  # parametrize input buffers: map each input buffer UOp to a PARAM with the correct slot index
+  linear = linear.substitute({u: UOp.param(i, u.dtype, u.shape, u.device) for i,u in enumerate(input_uops)}, walk=True)
+
+  # wrap SINKs with BEAM if jitbeam is set
+  if (jitbeam:=getenv("JITBEAM", BEAM.value)) >= 1: linear = graph_rewrite(linear, pm_beam, ctx=jitbeam, walk=True)
+
   linear = memory_plan_rewrite(linear, held_bufs)
-  if JIT < 2: linear = graph_split_rewrite(linear, set(input_buffers or []), max_batch_size=JIT_BATCH_SIZE.value)
+  if JIT < 2: linear = graph_split_rewrite(linear, max_batch_size=JIT_BATCH_SIZE.value)
   if VIZ: graph_rewrite(linear, PatternMatcher([]), name="View graphed linear")
-  return [ei.lower() for ei in linear_to_schedule(linear)]
+  return linear
 
 class GraphException(Exception): pass
 class JitError(Exception): pass
@@ -86,26 +102,20 @@ def _check_no_non_tensor_return(ret):
 
 def graph_class(dev): return dev.graph.func if isinstance(dev.graph, functools.partial) else dev.graph
 
-def get_input_replace(jit_cache: list[ExecItem], input_buffers:list[Buffer],
-                      orig_valid_positions: dict[int, set[int]]|None = None) -> dict[tuple[int, int], int]:
+def get_input_replace(jit_cache: list[ExecItem], input_buffers:list[Buffer]) -> dict[tuple[int, int], int]:
   input_replace: dict[tuple[int, int], int] = {}
   for j,ji in enumerate(jit_cache):
     for i,a in enumerate(ji.bufs):
-      if a in input_buffers:
-        # filter out positions that weren't valid inputs in the original capture (prevents aliasing bugs)
-        if orig_valid_positions is not None and i not in orig_valid_positions.get(id(ji), set()): continue
-        input_replace[(j,i)] = input_buffers.index(a)
+      if a in input_buffers: input_replace[(j,i)] = input_buffers.index(a)
   return input_replace
 
 class GraphRunner(Runner):
-  def __init__(self, linear:UOp|None, input_buffers:list[Buffer]|None,
-               jit_cache:list[ExecItem]|None=None, input_replace:dict[tuple[int,int],int]|None=None):
-    # TODO: captured jit as linear?
-    if linear is not None:
-      jit_cache = [ei.lower() for ei in linear_to_schedule(linear.src[0])]
-      for b in jit_cache_bufs(jit_cache): b.ensure_allocated()
-      input_replace = get_input_replace(jit_cache, input_buffers) if input_buffers else {}
-    self.jit_cache, self.input_replace = unwrap(jit_cache), input_replace or {}
+  def __init__(self, linear:UOp, input_buffers:list[Buffer]):
+    self.jit_cache = [ei.lower() for ei in linear_to_schedule(linear.src[0])]
+    for ei in self.jit_cache:
+      for b in ei.bufs:
+        if b is not None: b.ensure_allocated()
+    self.input_replace = get_input_replace(self.jit_cache, input_buffers) if input_buffers else {}
 
     self.var_vals_replace:dict[int, list[tuple[int, int]]] = {}
     self.launch_dims_replace:dict[int, tuple[int|None, int|None]] = {}
@@ -141,8 +151,6 @@ class GraphRunner(Runner):
     assert self.jit_cache[0].prg is not None
     super().__init__(colored(f"<batched {len(self.jit_cache)}>", "cyan"), self.jit_cache[0].prg.device.split(":")[0], estimates.simplify())
 
-  def __reduce__(self): return self.__class__, (None, None, self.jit_cache, self.input_replace)
-
   def updated_vars(self, var_vals: dict[str, int]):
     vals = [var_vals[v] for v in self.vars]
     for j, vidxs in self.var_vals_replace.items():
@@ -179,93 +187,55 @@ class GraphRunner(Runner):
 
   @staticmethod
   def supports_exec_item(batch_devs:list[Compiled], new_call:UOp) -> bool:
-    return new_call.src[0].op in (Ops.SINK, Ops.PROGRAM) and len(GraphRunner._all_devs(batch_devs, new_call)) == 1
+    return _unwrap_beam(new_call.src[0]).op in (Ops.SINK, Ops.PROGRAM) and len(GraphRunner._all_devs(batch_devs, new_call)) == 1
 
 # a marker for your graph supporting multiple devices of the same type
 class MultiGraphRunner(GraphRunner):
   @staticmethod
   def supports_exec_item(batch_devs:list[Compiled], new_call:UOp) -> bool:
     # Devices must be the same type
-    return new_call.src[0].op in (Ops.SINK, Ops.PROGRAM, Ops.COPY) and len(dedup([type(d) for d in GraphRunner._all_devs(batch_devs, new_call)])) == 1
-
-def get_out_buffers_for_ei(ei:ExecItem) -> list[Buffer]:
-  if isinstance(ei.prg, CompiledRunner): return [cast(Buffer, ei.bufs[out]) for out in ei.prg.p.outs if out not in ei.prg.p.ins]
-  if isinstance(ei.prg, (BufferCopy, BufferXfer, EncDec)): return [cast(Buffer, ei.bufs[0])]
-  if isinstance(ei.prg, GraphRunner): return dedup([b for inner in ei.prg.jit_cache for b in get_out_buffers_for_ei(inner)])
-  return []
-
-def update_depends(depends:set[Buffer|None], jit_cache:list[ExecItem]):
-  for ei in jit_cache:
-    if any(b in depends for b in ei.bufs): depends.update(get_out_buffers_for_ei(ei))
+    return _unwrap_beam(new_call.src[0]).op in (Ops.SINK, Ops.PROGRAM, Ops.COPY) \
+      and len(dedup([type(d) for d in GraphRunner._all_devs(batch_devs, new_call)])) == 1
 
 ReturnType = TypeVar('ReturnType')
 @dataclass
 class CapturedJit(Generic[ReturnType]):
   ret: Any  # includes the Tensors or any other returned object
-  jit_cache: list[ExecItem]
-  input_replace: dict[tuple[int, int], int]
-  extra_view_inputs: list[tuple[int, int, str, int, DType]]
+  linear: UOp
   expected_names: list[int|str]
   expected_input_info: list[tuple[UOp, tuple[Variable, ...], DType, str]]  # (view, variables, dtype, device) per input
 
-  def __reduce__(self):
-    # TODO: free_intermediates here?
-    return self.__class__, (self.ret, self.jit_cache, self.input_replace, self.extra_view_inputs, self.expected_names, self.expected_input_info)
+  def __reduce__(self): return self.__class__, (self.ret, self.linear, self.expected_names, self.expected_input_info)
 
-  def __post_init__(self):
-    self._jit_cache: list[ExecItem] = self.jit_cache
-    self._input_replace: dict[tuple[int, int], int] = self.input_replace
-    self._first_run = True
-    self._needs_rebuild = False
-    # precompute read-after-write hazard detection
-    self._output_to_writer = {b: j for j, ei in enumerate(self.jit_cache) for b in get_out_buffers_for_ei(ei)}
-    self._input_to_max_reader: dict[int, int] = {}
-    for (j, i), idx in self.input_replace.items():
-      # only buffers that were different during capture but alias at jit time (e.g. feeding output back as input) need the copy.
-      if self.jit_cache[j].bufs[i] not in get_out_buffers_for_ei(self.jit_cache[j]):
-        self._input_to_max_reader[idx] = max(self._input_to_max_reader.get(idx, -1), j)
-    self._clear_inputs()
+  @functools.cached_property
+  def _written_uops(self) -> set[UOp]:
+    out: set[UOp] = set()
+    for call in self.linear.toposort():
+      if call.op is not Ops.CALL: continue
+      non_bind = [s for s in call.src[1:] if s.op is not Ops.BIND]
+      outs, ins = _call_outs_ins(call)
+      out |= {non_bind[k] for k in outs - ins if non_bind[k].op in (Ops.BUFFER, Ops.BUFFER_VIEW)}
+    return out
 
-  def _clear_inputs(self):
-    for (j,i) in self._input_replace.keys(): self._jit_cache[j].bufs[i] = None
+  def __call__(self, input_uops:list[UOp], var_vals:dict[str, int]) -> ReturnType:
+    concrete = tuple(_copy_input(u) if u in self._written_uops else u for u in input_uops)
+    if DEBUG >= 1 and len(self.linear.src) >= 10: print(f"jit execs {len(self.linear.src)} calls")
+    run_linear(self.linear, var_vals, input_uops=concrete, jit=True)
+    return self.ret
 
   def free_intermediates(self):
-    depends: set[Buffer|None] = set([None])
-    update_depends(depends, self.jit_cache)
-    arenas = {b._base for b in depends if b is not None and b._base is not None}
-    to_free = {b for b in depends if b is not None} | {b for b in jit_cache_bufs(self.jit_cache) if b._base in arenas}
-    for b in to_free:
-      if hasattr(b, '_buf'): b.deallocate()
-    for a in arenas:
-      if a.allocated_views == 0 and a.is_allocated(): a.deallocate()
-    self.__post_init__()
-    self._needs_rebuild = True
-
-  # jit exec
-  def __call__(self, input_buffers:list[Buffer], var_vals:dict[str, int]) -> ReturnType:
-    # assign inputs
-    for idx, offset, device, size, dtype in self.extra_view_inputs:
-      input_buffers.append(Buffer(device, size, dtype, base=input_buffers[idx], offset=offset).ensure_allocated())
-
-    # copy aliased inputs to prevent read-after-write hazard
-    for i, ib in enumerate(input_buffers):
-      if (writer := self._output_to_writer.get(ib)) is not None and self._input_to_max_reader.get(i, -1) >= writer:
-        input_buffers[i] = Buffer(ib.device, ib.size, ib.dtype).ensure_allocated().copyin(ib.as_memoryview())
-
-    for (j,i),input_idx in self._input_replace.items(): self._jit_cache[j].bufs[i] = input_buffers[input_idx]
-
-    # allocate intermediates if freed on first run
-    if self._first_run:
-      for b in jit_cache_bufs(self.jit_cache): b.ensure_allocated()
-    if self._needs_rebuild:
-      for ei in self.jit_cache:
-        if isinstance(ei.prg, GraphRunner): ei.prg = type(ei.prg)(None, None, ei.prg.jit_cache, ei.prg.input_replace)
-    self._first_run = self._needs_rebuild = False
-
-    if DEBUG >= 1 and len(self._jit_cache) >= 10: print(f"jit execs {len(self._jit_cache)} kernels")
-    for ei in self._jit_cache: ei.run(var_vals, jit=True)
-    self._clear_inputs()
-    return self.ret
+    # drop graph runners
+    for call in self.linear.src:
+      if call.src[0].op is Ops.CUSTOM_FUNCTION and call.src[0].arg == "graph": graph_cache.pop(call.src[0], None)
+    bases: set[Buffer] = set()
+    for u in self._written_uops:
+      try: buf = u.buffer
+      except Exception: continue
+      for b in (buf.bufs if isinstance(buf, MultiBuffer) else [buf]):
+        if hasattr(b, '_buf'): b.deallocate()
+        if b._base is not None: bases.add(b._base)
+    for a in bases:
+      if a.is_allocated() and a.allocated_views == 0: a.deallocate()
 
 def _prepare_jit_inputs(args, kwargs):
   input_tensors: list[tuple[int|str, Tensor]] = [(name,t) for name,t in list(enumerate(args))+sorted(kwargs.items()) if t.__class__ is Tensor]
@@ -278,13 +248,14 @@ def _prepare_jit_inputs(args, kwargs):
   input_uops: list[UOp] = flatten([t.uop.src if t.uop.op is Ops.MULTI else [t.uop] for t in tensors])
   if any(u.base.op is Ops.CONST for u in input_uops):
     raise JitError("JIT inputs cannot be const, create a buffer with .contiguous()")
-  input_buffers: list[Buffer] = flatten([b.bufs if isinstance(b, MultiBuffer) else [b] for u in input_uops if (b:=u.base.realized) is not None])
-  if len(set(input_buffers)) != len(input_buffers): raise JitError("duplicate inputs to JIT")
+  # collect buffer UOps (including MultiBuffer)
+  input_buf_uops: list[UOp] = [u.base for u in input_uops if u.base.realized is not None]
+  if len(set(input_buf_uops)) != len(input_buf_uops): raise JitError("duplicate inputs to JIT")
   inputs = [(*(u.substitute({u.base:UOp(Ops.NOOP)}, extra_pm=mop_cleanup).unbind_all()), u.dtype, u.device) for u in input_uops]
   _var_vals = merge_dicts([x[1] for x in inputs] + [dict(v.unbind() for v in (args + tuple(kwargs.values())) if isinstance(v, UOp))])
   var_vals = {k.expr:v for k,v in _var_vals.items()}
   expected_input_info = [(x[0], tuple(sorted(x[1].keys(), key=lambda v: v.expr)), x[2], x[3]) for x in inputs]
-  return input_buffers, var_vals, names, expected_input_info
+  return input_buf_uops, var_vals, names, expected_input_info
 
 class TinyJit(Generic[ReturnType]):
   def __init__(self, fxn:Callable[..., ReturnType]|None, captured:CapturedJit|None=None, prune=False):
@@ -305,16 +276,10 @@ class TinyJit(Generic[ReturnType]):
     assert self.captured is not None, "can't pickle an uncaptured JIT"
     return self.__class__, (None, self.captured)
 
-  # keep legacy code working
-  @property
-  def jit_cache(self) -> list[ExecItem]: return self.captured._jit_cache if self.captured is not None else []
-  @property
-  def input_replace(self) -> dict[tuple[int, int], int]: return self.captured._input_replace if self.captured is not None else {}
-
   def __get__(self, obj, objtype): return functools.partial(self.__call__, obj) # add support for instance methods
 
   def __call__(self, *args, **kwargs) -> ReturnType:
-    input_buffers, var_vals, names, expected_input_info = _prepare_jit_inputs(args, kwargs)
+    input_buf_uops, var_vals, names, expected_input_info = _prepare_jit_inputs(args, kwargs)
     if not JIT or self.cnt == 0:
       # jit ignore
       assert self.fxn is not None
@@ -333,46 +298,28 @@ class TinyJit(Generic[ReturnType]):
       finally: capturing.clear()
       if not len(self._linears): raise JitError("didn't JIT anything!")
       _check_no_non_tensor_return(ret)
-      if DEBUG >= 1: print(f"JIT captured {len(self._linears)} linears with {len(input_buffers)} inputs")
+      if DEBUG >= 1: print(f"JIT captured {len(self._linears)} linears with {len(input_buf_uops)} inputs")
 
-      # combine all captured linears into one, memory plan, and convert to ExecItems
+      # combine all captured linears into one, memory plan, and graph split
       big_linear = UOp(Ops.LINEAR, src=tuple(flatten([l.src for l in self._linears])))
       del self._linears
 
       if self.prune:
-        big_linear, onetime_linear = prune_linear(big_linear, {k for k,v in buffers.items() if isinstance(v, Buffer) and v in set(input_buffers)})
+        big_linear, onetime_linear = prune_linear(big_linear, set(input_buf_uops))
         if DEBUG >= 1: print(f"pruned from {len(big_linear.src) + len(onetime_linear.src)} -> {len(big_linear.src)} kernels")
-        for ei in (si.lower() for si in linear_to_schedule(onetime_linear)):
-          for b in ei.bufs: cast(Buffer, b).ensure_allocated()
-          ei.run(var_vals, jit=True)
+        run_linear(onetime_linear, var_vals)
 
       held_bufs = set(buffers) | {t.uop.buf_uop for t in get_parameters(ret) if t.uop.buf_uop.op is Ops.BUFFER}
-      with Context(BEAM=getenv("JITBEAM", BEAM.value)):
-        jit_cache = jit_lower(big_linear, held_bufs, input_buffers)
-
-      # track inputs that are views of buffers
-      # TODO: eventually expected_buffers should live in ExecItem
-      extra_view_inputs: list[tuple[int, int, str, int, DType]] = []
-      for item in jit_cache:
-        for b in item.bufs:
-          if b is not None and b._base is not None and b._base in input_buffers:
-            input_buffers.append(b)
-            extra_view_inputs.append((input_buffers.index(b.base), b.offset, b.device, b.size, b.dtype))
-
-      input_replace = get_input_replace(jit_cache, input_buffers)
-      if DEBUG >= 1 and len(set(input_replace.values())) != len(input_buffers): print("WARNING: some input tensors not found")
-
-      # exec
-      for ei in jit_cache: ei.run(var_vals)
-
-      self.captured = CapturedJit(ret, jit_cache, input_replace, extra_view_inputs, names, expected_input_info)
+      linear = jit_lower(big_linear, held_bufs, input_buf_uops)
+      self.captured = CapturedJit(ret, linear, names, expected_input_info)
+      ret = self.captured(input_buf_uops, var_vals)
     elif self.cnt >= 2:
       # jit exec
       assert self.captured is not None
       if self.captured.expected_names != names: raise JitError(f"args mismatch in JIT: {self.captured.expected_names=} != {names}")
       if self.captured.expected_input_info != expected_input_info:
         raise JitError(f"args mismatch in JIT: {self.captured.expected_input_info=} != {expected_input_info=}")
-      ret = self.captured(input_buffers, var_vals)
+      ret = self.captured(input_buf_uops, var_vals)
 
     self.cnt += 1
     return ret

tinygrad/engine/realize.py

@@ -1,11 +1,11 @@
-from typing import cast, Callable
-import time, pprint, random, itertools, math
+from typing import cast, Callable, Iterator
+import time, pprint, random, itertools, math, contextlib, weakref
 from dataclasses import dataclass, replace, field
 from tinygrad.helpers import all_same, colored, DEBUG, GlobalCounters, ansilen, NOOPT, all_int, Metadata, TRACEMETA, TracingKey
-from tinygrad.helpers import DEVECTORIZE, time_to_str, VALIDATE_WITH_CPU, cpu_profile, PROFILE, ProfilePointEvent, cpu_events, prod, unwrap
-from tinygrad.helpers import EMULATED_DTYPES
-from tinygrad.uop.ops import Ops, PatternMatcher, UOp, UPat, sym_infer
-from tinygrad.device import Device, Buffer
+from tinygrad.helpers import BEAM, DEVECTORIZE, size_to_str, time_to_str, VALIDATE_WITH_CPU, cpu_profile, PROFILE, ProfilePointEvent, cpu_events
+from tinygrad.helpers import prod, unwrap, EMULATED_DTYPES, flatten
+from tinygrad.uop.ops import Ops, PatternMatcher, UOp, UPat, sym_infer, buffers, graph_rewrite
+from tinygrad.device import Device, Buffer, MultiBuffer
 from tinygrad.renderer import ProgramSpec, Estimates
 from tinygrad.codegen import get_program
 
@@ -216,3 +216,110 @@ def run_schedule(schedule:list[ExecItem], var_vals:dict[str, int]|None=None, do_
       np.testing.assert_allclose(bufs[0].numpy(), nb[0].numpy(), rtol=1e-3, atol=1e-3)
     else:
       ei.run(var_vals, do_update_stats=do_update_stats)
+
+# **************** run linear ****************
+
+@dataclass
+class ExecContext:
+  var_vals: dict[str, int] = field(default_factory=dict)
+  input_uops: tuple[UOp, ...] = ()
+  do_update_stats: bool = True
+  jit: bool = False
+
+def _resolve(b:UOp, inputs:tuple[UOp, ...]) -> UOp:
+  if b.op is Ops.BUFFER_VIEW and b.src[0].op is Ops.PARAM: return b.replace(src=(inputs[b.src[0].arg], *b.src[1:]))
+  return inputs[b.arg] if b.op is Ops.PARAM else b
+def resolve_params(ctx:ExecContext, call:UOp) -> list[UOp]: return [_resolve(b, ctx.input_uops) for b in call.src[1:] if b.op is not Ops.BIND]
+
+@contextlib.contextmanager
+def track_stats(ctx:ExecContext, call:UOp, device:str, display_name:str, estimates:Estimates, bufs:list[Buffer], var_vals:dict[str, int],
+                outputs=(0,), inputs=(1,), first_run=False):
+  if PROFILE: cpu_events.append(ProfilePointEvent(device, "exec", len(cpu_events), {"metadata": call.arg.metadata, "var_vals": var_vals,
+                                                  "bufs": [b.trace_num for b in bufs], "name": display_name, "outputs": outputs, "inputs": inputs}))
+  timing: list[float|None] = [None]
+  if DEBUG >= 2: st = time.perf_counter()
+  yield timing
+  if not ctx.do_update_stats: return
+  if DEBUG >= 2 and timing[0] is None:
+    Device[device].synchronize()
+    timing[0] = time.perf_counter() - st
+  update_stats(display_name, device, estimates, var_vals, timing[0], len(bufs), jit=ctx.jit, metadata=call.arg.metadata, first_run=first_run)
+
+def unwrap_multi(call:UOp, resolved:list[UOp]) -> Iterator[tuple[list[Buffer], dict[str, int]]]:
+  bufs = [b.buffer for b in resolved]
+  if not any(isinstance(b, MultiBuffer) for b in bufs): yield cast(list[Buffer], bufs), {}
+  else:
+    dnum = next((x.expr for x in call.src[0].variables() if x.expr == '_device_num'), None)
+    for j, per_dev in enumerate(zip(*[cast(MultiBuffer, b).bufs for b in bufs])): yield list(per_dev), {dnum: j} if dnum else {}
+
+def exec_view(ctx:ExecContext, call, ast):
+  resolved = resolve_params(ctx, call)
+  bufs = [cast(Buffer, b.buffer) for b in resolved]
+  bv = bufs[1].view(resolved[0].arg, ast.dtype, ast.arg[1]*bufs[1].dtype.itemsize)
+  with track_stats(ctx, call, bv.device, colored(f"view {bv.nbytes:8d} @ {bv.offset:<10d}", "yellow"), Estimates(), [bv, bufs[1]], ctx.var_vals):
+    buffers[resolved[0]] = bv
+
+def exec_copy(ctx:ExecContext, call, ast):
+  for bufs, device_vars in unwrap_multi(call, resolve_params(ctx, call)):
+    dest, src = bufs[0].ensure_allocated(), bufs[1].ensure_allocated()
+    xfer = hasattr(alc:=Device[dest.device].allocator,'_transfer') and alc.supports_transfer and dest.device.split(":")[0]==src.device.split(":")[0]
+    prg = (BufferXfer if xfer else BufferCopy)(dest.nbytes, dest.device, src.device)
+    with track_stats(ctx, call, dest.device, prg.display_name, Estimates(lds=dest.nbytes, mem=dest.nbytes), [dest, src], ctx.var_vals):
+      prg.copy(dest, src)
+
+def exec_kernel(ctx:ExecContext, call, ast):
+  sink = ast.src[0] if ast.op is Ops.BEAM else ast
+
+  for bufs, device_vars in unwrap_multi(call, resolve_params(ctx, call)):
+    var_vals = {**ctx.var_vals, **device_vars}
+    prg = get_runner(bufs[0].device, ast)
+    prg_bufs = [bufs[i].ensure_allocated() for i in prg.p.globals]
+
+    if VALIDATE_WITH_CPU and sink.op is Ops.SINK:
+      cpu_bufs = [Buffer("CPU", b.size, b.dtype).ensure_allocated().copyin(b.ensure_allocated().as_memoryview()) for b in bufs]
+
+    with track_stats(ctx, call, prg.device, prg.display_name, prg.estimates, prg_bufs, var_vals,
+                     outputs=tuple(prg.p.outs), inputs=tuple(prg.p.ins), first_run=prg.first_run) as timing:
+      timing[0] = prg(prg_bufs, var_vals, wait=DEBUG >= 2)
+      prg.first_run = False
+
+    if VALIDATE_WITH_CPU and sink.op is Ops.SINK:
+      import numpy as np
+      cpu_prg = get_runner("CPU", sink)
+      cpu_prg([cpu_bufs[i] for i in cpu_prg.p.globals], var_vals, wait=False)
+      for i in prg.p.outs: np.testing.assert_allclose(prg_bufs[i].numpy(), cpu_bufs[i].numpy(), rtol=1e-3, atol=1e-3)
+
+def exec_encdec(ctx:ExecContext, call, ast):
+  bufs = [cast(Buffer, b.buffer).ensure_allocated() for b in resolve_params(ctx, call)]
+  shape, pos_var = tuple(s.arg for s in ast.src if s.op is Ops.CONST), ast.variables()[0].expr
+  with track_stats(ctx, call, bufs[0].device, colored(f"enc/dec {size_to_str(bufs[0].nbytes)}", "yellow"),
+                   Estimates(lds=bufs[0].nbytes, mem=bufs[0].nbytes), bufs, ctx.var_vals):
+    bufs[0].allocator._encode_decode(bufs[0]._buf, bufs[1]._buf, bufs[2]._buf, [x._buf for x in bufs[3:]], shape, ctx.var_vals[pos_var])
+
+graph_cache:weakref.WeakKeyDictionary[UOp, Runner] = weakref.WeakKeyDictionary()
+def exec_graph(ctx:ExecContext, call, cf):
+  inputs = resolve_params(ctx, call)
+  bufs = flatten([b.bufs if isinstance(b, MultiBuffer) else [b] for b in (u.buffer for u in inputs)])
+  if (runner:=graph_cache.get(cf)) is None:
+    sub = cf.substitute(dict(zip(cf.src[1:], inputs)))
+    graph_cache[cf] = runner = Device[cf.device if isinstance(cf.device, str) else cf.device[0]].graph(sub, bufs)
+  with track_stats(ctx, call, runner.device, runner.display_name, runner.estimates, bufs, ctx.var_vals) as t:
+    t[0] = runner(bufs, ctx.var_vals, wait=DEBUG >= 2)
+
+pm_beam = PatternMatcher([
+  (UPat(Ops.CALL, src=(UPat(Ops.SINK, name="sink"),), name="call", allow_any_len=True),
+   lambda ctx,call,sink: call.replace(src=(UOp(Ops.BEAM, src=(sink,), arg=ctx), *call.src[1:]))),
+])
+
+pm_exec = PatternMatcher([
+  (UPat(Ops.CALL, src=(UPat(Ops.BUFFER_VIEW, name="ast"),), name="call", allow_any_len=True), exec_view),
+  (UPat(Ops.CALL, src=(UPat(Ops.COPY, name="ast"),), name="call", allow_any_len=True), exec_copy),
+  (UPat(Ops.CALL, src=(UPat((Ops.SINK, Ops.PROGRAM, Ops.BEAM), name="ast"),), name="call", allow_any_len=True), exec_kernel),
+  (UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION, arg="encdec", name="ast"),), name="call", allow_any_len=True), exec_encdec),
+  (UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION, arg="graph", name="cf"),), name="call", allow_any_len=True), exec_graph),
+])
+
+def run_linear(linear:UOp, var_vals:dict[str, int]|None=None, input_uops:tuple[UOp, ...]=(), do_update_stats=True, jit=False):
+  if BEAM >= 1: linear = graph_rewrite(linear, pm_beam, ctx=BEAM.value, name="add beam")
+  ctx = ExecContext(var_vals or {}, input_uops, do_update_stats, jit)
+  for call in linear.src: pm_exec.rewrite(call, ctx)

tinygrad/function.py

@@ -70,7 +70,7 @@ class _function(Generic[ReturnType]):
     if not self.allow_implicit:
       implicit_buffers = [x for x in call_uops[num_explicit:] if x.op is Ops.BUFFER]
       if implicit_buffers:
-        buf_strs = '\n  '.join(f"{i}: dtype={b.dtype}, size={b.size}, device={b.device}" for i,b in enumerate(implicit_buffers))
+        buf_strs = '\n  '.join(f"{i}: dtype={b.dtype}, size={b.arg}, device={b.device}" for i,b in enumerate(implicit_buffers))
         raise RuntimeError(f"function {name} has {len(implicit_buffers)} implicit buffer(s), but allow_implicit=False\n  {buf_strs}")
 
     # assign output

tinygrad/gradient.py

@@ -93,17 +93,18 @@ def compute_gradient(root:UOp, root_grad:UOp, targets:set[UOp]) -> dict[UOp, UOp
   grads: dict[UOp, UOp] = {root: root_grad}
   for t0 in reversed(walk):
     if t0 not in grads or grads[t0].op is Ops.NOOP: continue
-    # GETTUPLE: accumulate gradient into a TUPLE UOp on the CALL, process when we hit the CALL
+    # GETTUPLE: accumulate gradient into a TUPLE UOp on the FUNCTION, process when we hit the FUNCTION
     if t0.op is Ops.GETTUPLE:
-      k = t0.src[0]  # the CALL
-      assert k.op is Ops.CALL and k.src[0].op is Ops.TUPLE
+      k = t0.src[0]  # the FUNCTION
+      assert k.op is Ops.FUNCTION and k.src[0].op is Ops.TUPLE
       n_outputs = len(k.src[0].src)
       prev = grads[k].src if k in grads else tuple(UOp(Ops.NOOP) for _ in range(n_outputs))
       grads[k] = UOp.maketuple(*(prev[i] + grads[t0] if i == t0.arg and prev[i].op is not Ops.NOOP else
                                  grads[t0] if i == t0.arg else prev[i] for i in range(n_outputs)))
       continue
-    # CALL: pass needed param set so backward only computes required gradients
-    if t0.op is Ops.CALL:
+    # FUNCTION/CALL: pass needed param set so backward only computes required gradients
+    # (FUNCTION uses implicit TUPLE gradient or grad_fxn; CALL requires an explicit grad_fxn)
+    if t0.op in {Ops.FUNCTION, Ops.CALL}:
       needed = {i for i, arg in enumerate(t0.src[1:]) if arg in targets or in_target_path.get(arg, False)}
       lgrads:tuple[UOp|None, ...]|None = call_gradient(grads[t0], t0, needed)
     else:

tinygrad/helpers.py

@@ -32,6 +32,7 @@ def argsort(x): return type(x)(sorted(range(len(x)), key=x.__getitem__))
 def all_same(items:Sequence): return all(x == items[0] for x in items)  # works for empty input
 def all_int(t: Sequence[Any]) -> TypeGuard[tuple[int, ...]]: return all(isinstance(s, int) for s in t)
 def colored(st, color:str|None, background=False): # replace the termcolor library
+  if NO_COLOR: return st
   colors = ['black', 'red', 'green', 'yellow', 'blue', 'magenta', 'cyan', 'white']
   return f"\u001b[{10*background+60*(color.upper() == color)+30+colors.index(color.lower())}m{st}\u001b[0m" if color is not None else st
 def colorize_float(x: float): return colored(f"{x:7.2f}x", 'green' if x < 0.75 else 'red' if x > 1.15 else 'yellow')
@@ -44,6 +45,12 @@ def flatten(l:Iterable[Iterable[T]]): return [item for sublist in l for item in
 def fully_flatten(l):
   if not (hasattr(l, "__len__") and hasattr(l, "__getitem__")) or isinstance(l, str): return [l]
   return [l[()]] if hasattr(l, "shape") and l.shape == () else [x for li in l for x in fully_flatten(li)]
+#  `(padding_left, padding_right, padding_top, padding_bottom, ...)` ->  `(..., (padding_top, padding_bottom), (padding_left, padding_right))`
+def flat_to_grouped(padding:Sequence[T]) -> tuple[tuple[T, T], ...]: return tuple(zip(padding[-2::-2], padding[::-2]))
+def resolve_pool_pads(padding:int|Sequence[int], dims:int) -> Sequence[int]:
+  if not isinstance(padding, int) and not (len(padding) == 2*dims or len(padding) == dims):
+    raise ValueError(f"Padding must be an int or a sequence of length {dims} or {2*dims}, but got {padding=} with {dims=}.")
+  return [padding]*2*dims if isinstance(padding, int) else (padding if len(padding) == 2*dims else [p for p in padding for _ in range(2)][::-1])
 def fromimport(mod, frm): return getattr(__import__(mod, fromlist=[frm]), frm)
 def _is_balanced(s:str) -> bool: return (d := 0, all((d := d + (c == '(') - (c == ')')) >= 0 for c in s))[1] and d == 0
 def strip_parens(fst:str) -> str: return fst[1:-1] if fst[:1]=='(' and fst[-1:]==')' and _is_balanced(fst[1:-1]) else fst
@@ -224,7 +231,7 @@ class _DEV(ContextVar):
 DEV, DEBUG, BEAM, NOOPT = _DEV("DEV", ""), ContextVar("DEBUG", 0), ContextVar("BEAM", 0), ContextVar("NOOPT", 0)
 IMAGE, FLOAT16, OPENPILOT_HACKS = ContextVar("IMAGE", 0), ContextVar("FLOAT16", 0), ContextVar("OPENPILOT_HACKS", 0)
 JIT, JIT_BATCH_SIZE = ContextVar("JIT", 2 if OSX and ARCH_X86 else 1), ContextVar("JIT_BATCH_SIZE", 32)
-WINO, CAPTURING, TRACEMETA = ContextVar("WINO", 0), ContextVar("CAPTURING", 1), ContextVar("TRACEMETA", 1)
+WINO, CAPTURING, TRACEMETA, NO_COLOR = ContextVar("WINO", 0), ContextVar("CAPTURING", 1), ContextVar("TRACEMETA", 1), ContextVar("NO_COLOR", 0)
 USE_TC, TC_SELECT, TC_OPT, AMX = ContextVar("TC", 1), ContextVar("TC_SELECT", -1), ContextVar("TC_OPT", 0), ContextVar("AMX", 0)
 TRANSCENDENTAL, NOLOCALS = ContextVar("TRANSCENDENTAL", 1), ContextVar("NOLOCALS", 0)
 SPLIT_REDUCEOP, NO_MEMORY_PLANNER, LRU = ContextVar("SPLIT_REDUCEOP", 1), ContextVar("NO_MEMORY_PLANNER", 0), ContextVar("LRU", 1)
@@ -472,14 +479,14 @@ def cpu_objdump(lib, objdump_tool='objdump'):
     pathlib.Path(f.name).write_bytes(lib)
     print(system(f"{objdump_tool} -d {f.name}"))
 
-def capstone_flatdump(lib: bytes):
+def capstone_flatdump(lib: bytes, arch:str):
   try: import capstone
   except ImportError:
     print("Disassembler Error: Capstone not installed.")
     return
-  match platform.machine():
-    case 'x86_64' | 'AMD64': cs = capstone.Cs(capstone.CS_ARCH_X86, capstone.CS_MODE_64)
-    case 'aarch64' | 'arm64': cs = capstone.Cs(capstone.CS_ARCH_ARM64, capstone.CS_MODE_ARM)
+  match arch:
+    case 'x86_64': cs = capstone.Cs(capstone.CS_ARCH_X86, capstone.CS_MODE_64)
+    case 'arm64': cs = capstone.Cs(capstone.CS_ARCH_ARM64, capstone.CS_MODE_ARM)
     case machine: raise NotImplementedError(f"Capstone disassembly isn't supported for {machine}")
   cs.skipdata = True
   for instr in cs.disasm(lib, 0):

tinygrad/llm/__main__.py

@@ -0,0 +1,2 @@
+from tinygrad.llm.cli import main
+if __name__ == "__main__": main()