recursive stuff works

works but bad
add callify method
2026-08-15 00:18:27 +00:00 · 2026-02-24 15:15:36 +08:00 · 2026-02-24 14:40:39 +08:00 · 2026-02-24 11:44:33 +08:00
109 changed files with 11293 additions and 3136 deletions
@@ -233,7 +233,7 @@ runs:
      shell: bash
      run: |
        sudo mkdir -p /usr/local/lib
-        curl -s -H "Authorization: token $GH_TOKEN" curl -s https://api.github.com/repos/tinygrad/amdcomgr_dylib/releases/latest | \
+        curl -s -H "Authorization: token $GH_TOKEN" curl -s https://api.github.com/repos/nimlgen/amdcomgr_dylib/releases/latest | \
          jq -r '.assets[] | select(.name == "libamd_comgr.dylib").browser_download_url' | \
          sudo xargs curl -fL -o /usr/local/lib/libamd_comgr.dylib
        cargo build --release --manifest-path ./extra/remu/Cargo.toml
@@ -32,7 +32,6 @@ jobs:
    - name: Setup Environment
      uses: ./.github/actions/setup-tinygrad
      with:
-        key: 'autogen'
        opencl: 'true'
        amd: 'true'
        cuda: 'true'
@@ -82,7 +81,6 @@ jobs:
    - name: Setup Environment
      uses: ./.github/actions/setup-tinygrad
      with:
-        key: 'autogen-mac'
        llvm: 'true'
    - name: Regenerate autogen files
      run: |
@@ -112,8 +110,6 @@ jobs:
      uses: actions/checkout@v4
    - name: Setup Environment
      uses: ./.github/actions/setup-tinygrad
-      with:
-        key: 'autogen-comgr'
    - name: Install autogen support packages
      run: |
        wget https://repo.radeon.com/rocm/rocm.gpg.key -O - | gpg --dearmor | sudo tee /etc/apt/keyrings/rocm.gpg > /dev/null
@@ -520,7 +520,7 @@ jobs:
      run: time BENCHMARK_LOG=cifar AMD=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
    - name: Run full CIFAR training steps w 6 GPUS
      run: time BENCHMARK_LOG=cifar_6gpu AMD=1 DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
-    # TODO: broken on some of the machines
+    # this needs to be mocked and testable on a local machine
    #- name: Test full tinyfs load
    #  run: TINYFS_ENDPOINT=10.0.52.11:6767 PYTHONPATH=. python extra/tinyfs/fetch_file.py --hash d734f5e3be9f1e9d863bfaa4fc6c1ef2 --len 175866113 --dest mapping.json --check
    - name: Run process replay tests
@@ -617,27 +617,6 @@ jobs:
    - name: Run process replay tests
      run: cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && PYTHONPATH=. python3 process_replay.py

-  testcommausbgpubenchmark:
-    name: UsbGPU Benchmark (comma)
-    runs-on: [self-hosted, Linux, comma4]
-    timeout-minutes: 20
-    defaults:
-      run:
-        shell: bash -e -o pipefail {0}
-    if: github.repository_owner == 'tinygrad'
-    steps:
-    - name: Checkout Code
-      uses: actions/checkout@v4
-    - name: setup staging db
-      if: github.ref == 'refs/heads/update_benchmark_staging'
-      run: |
-        echo "CACHEDB=/tmp/staging.db" >> $GITHUB_ENV
-        rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
-    - name: openpilot compile3 0.10.1 driving_vision
-      run: BENCHMARK_LOG=usbgpu_openpilot_0_10_1_vision PYTHONPATH="." DEV=AMD AMD_LLVM=1 AMD_IFACE=USB ASSERT_MIN_STEP_TIME=50 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
-    - name: openpilot load_pickle 0.10.1 driving_vision
-      run: BENCHMARK_LOG=usbgpu_openpilot_0_10_1_vision_load_pickle PYTHONPATH="." DEV=AMD AMD_IFACE=USB ASSERT_MIN_LOAD_TIME=15 python3 examples/openpilot/load_pickle.py
-
  testreddriverbenchmark:
    name: AM Benchmark
    runs-on: [self-hosted, Linux, tinyboxrandom]
@@ -244,37 +244,6 @@ jobs:
    - name: Run TYPED=1
      run: CHECK_OOB=0 DEV=CPU TYPED=1 python test/test_tiny.py

-  nulltest:
-    name: Null Tests
-    runs-on: ubuntu-latest
-    timeout-minutes: 15
-
-    steps:
-    - name: Checkout Code
-      uses: actions/checkout@v4
-    - name: Setup Environment
-      uses: ./.github/actions/setup-tinygrad
-      with:
-        key: unittest-13
-        pydeps: "pillow ftfy regex pre-commit"
-        deps: testing_unit
-        llvm: 'true'
-        amd: 'true'
-    - name: Run NULL backend tests
-      run: NULL=1 python -m pytest -n=auto test/null/ --durations=20
-    - name: Run targetted tests on NULL backend
-      run: NULL=1 python3 -m unittest test.backend.test_multitensor.TestMultiTensor.test_data_parallel_resnet_train_step
-    # TODO: too slow
-    # - name: Run SDXL on NULL backend
-    #   run: NULL=1 DEBUG=1 python3 examples/sdxl.py --seed 0 --noshow --timing --fakeweights
-    - name: Run Clip tests for SD MLPerf on NULL backend
-      run: NULL=1 python -m pytest -n=auto test/external/mlperf_stable_diffusion/external_test_models.py::TestOpenClip --durations=20
-    - name: Run AMD emulated BERT training on NULL backend
-      run: EMULATE=AMD_RDNA4 NULL=1 NULL_ALLOW_COPYOUT=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
-    # TODO: support fake weights
-    #- name: Run LLaMA 7B on 4 fake devices
-    #  run: NULL=1 python3 examples/llama.py --gen 1 --size 7B --shard 4 --prompt "Hello." --count 3 --temperature 0 --timing
-
  unittest:
    name: Unit Tests
    runs-on: ubuntu-latest
@@ -299,6 +268,20 @@ jobs:
      run: |
        CPU=1 python test/null/test_device.py TestRunAsModule.test_module_runs
        CPU=1 python -m pytest -n=auto test/unit/ --durations=20
+    - name: Run NULL backend tests
+      run: NULL=1 python -m pytest -n=auto test/null/ --durations=20
+    - name: Run targetted tests on NULL backend
+      run: NULL=1 python3 -m unittest test.backend.test_multitensor.TestMultiTensor.test_data_parallel_resnet_train_step
+    # TODO: too slow
+    # - name: Run SDXL on NULL backend
+    #   run: NULL=1 DEBUG=1 python3 examples/sdxl.py --seed 0 --noshow --timing --fakeweights
+    - name: Run Clip tests for SD MLPerf on NULL backend
+      run: NULL=1 python -m pytest -n=auto test/external/mlperf_stable_diffusion/external_test_models.py::TestOpenClip --durations=20
+    - name: Run AMD emulated BERT training on NULL backend
+      run: EMULATE=AMD_RDNA4 NULL=1 NULL_ALLOW_COPYOUT=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
+    # TODO: support fake weights
+    #- name: Run LLaMA 7B on 4 fake devices
+    #  run: NULL=1 python3 examples/llama.py --gen 1 --size 7B --shard 4 --prompt "Hello." --count 3 --temperature 0 --timing
    - name: Run GC tests
      run: python test/external/external_uop_gc.py
    - name: External Benchmark Schedule
@@ -661,7 +644,7 @@ jobs:
          sudo apt-get update
          sudo apt-get install llvm-21 llvm-21-tools cloc
      - name: Install rocprof-trace-decoder
-        run: sudo PYTHONPATH="." ./extra/sqtt/install_rocprof_decoder.py
+        run: sudo PYTHONPATH="." ./extra/sqtt/install_sqtt_decoder.py
      - name: Run AMD renderer tests
        run: AMD_LLVM=0 python -m pytest -n=auto test/amd/ --durations 20
      - name: Run AMD renderer tests (AMD_LLVM=1)
@@ -396,7 +396,6 @@ def batch_load_retinanet(dataset, val:bool, base_dir:Path, batch_size:int=32, sh
      queue_in.put((idx, img, tgt))

  def _setup_shared_mem(shm_name:str, size:tuple[int, ...], dtype:dtypes) -> tuple[shared_memory.SharedMemory, Tensor]:
-    shm_name = f"{shm_name}_{os.getpid()}"
    if os.path.exists(f"/dev/shm/{shm_name}"): os.unlink(f"/dev/shm/{shm_name}")
    shm = shared_memory.SharedMemory(name=shm_name, create=True, size=prod(size))
    shm_tensor = Tensor.empty(*size, dtype=dtype, device=f"disk:/dev/shm/{shm_name}")
@@ -3,7 +3,7 @@ from pathlib import Path
 import multiprocessing

 from tinygrad import Device, GlobalCounters, Tensor, TinyJit, dtypes
-from tinygrad.helpers import getenv, BEAM, WINO, round_up, diskcache_clear, Profiling, profile_marker, DEBUG
+from tinygrad.helpers import getenv, BEAM, WINO, round_up, diskcache_clear, Profiling, profile_marker
 from tinygrad.nn.state import get_parameters, get_state_dict, load_state_dict, safe_load, safe_save
 from tinygrad.nn.optim import LAMB, LARS, SGD, OptimizerGroup, Adam, AdamW

@@ -1335,14 +1335,9 @@ def train_llama3():
  model_params = MODEL_PARAMS[getenv("LLAMA3_SIZE", "8B")]["args"]
  # vocab_size from the mixtral tokenizer
  if not SMALL: model_params |= {"vocab_size": 32000}
-  real_vocab_size = model_params['vocab_size']
  if (llama_layers:=getenv("LLAMA_LAYERS")) != 0: model_params['n_layers'] = llama_layers
  print(f"model parameters: {model_params}")

-  # pad vocab
-  if (MP := getenv("MP", 1)) > 1: model_params['vocab_size'] = round_up(model_params['vocab_size'], 256 * MP)
-  vocab_mask:Tensor = Tensor.arange(model_params['vocab_size']).reshape(1, 1, -1) >= real_vocab_size
-
  model = Transformer(**model_params, max_context=SEQLEN, jit=False, disable_kv_cache=True)
  params = get_parameters(model)
  # weights are all bfloat16 for now
@@ -1357,8 +1352,6 @@ def train_llama3():
    for v in get_parameters(model):
      v.shard_(device, axis=None)

-    vocab_mask.shard_(device, axis=None)
-
  if (MP := getenv("MP", 1)) > 1:
    device = tuple(f"{Device.DEFAULT}:{i}" for i in range(MP))
    for k,v in get_state_dict(model).items():
@@ -1366,7 +1359,6 @@ def train_llama3():
      elif '.attention.wq' in k: v.shard_(device, axis=0)
      elif '.attention.wk' in k: v.shard_(device, axis=0)
      elif '.attention.wv' in k: v.shard_(device, axis=0)
-      elif '.attention.wqkv' in k: v.shard_(device, axis=0)
      elif '.attention.wo' in k: v.shard_(device, axis=1)
      elif '.feed_forward.w1.' in k: v.shard_(device, axis=0)
      elif '.feed_forward.w2.' in k: v.shard_(device, axis=1)
@@ -1379,18 +1371,14 @@ def train_llama3():
      # prevents memory spike on device 0
      v.realize()

-    vocab_mask.shard_(device, axis=2).realize()
-
  optim_device = "CPU" if getenv("OFFLOAD_OPTIM") else None
  optim = GradAccClipAdamW(get_parameters(model), 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 = p.empty_like().realize()
+    p.grad = p.zeros_like().contiguous().realize()
  grads: list[Tensor] = [p.grad for p in optim.params]
-  for p in optim.params:
-    p.grad.assign(p.grad.zeros_like()).realize()

  scheduler = CosineAnnealingLRWithWarmup(optim, opt_base_learning_rate, opt_end_learning_rate, opt_learning_rate_warmup_steps, opt_learning_rate_decay_steps)

@@ -1405,15 +1393,15 @@ def train_llama3():

  @TinyJit
  def minibatch(tokens:Tensor):
+    tokens = tokens.to(None)
    if (DP := getenv("DP", 1)) > 1:
      device = tuple(f"{Device.DEFAULT}:{i}" for i in range(DP))
-      tokens = tokens.to(None).shard(device, 0)
+      tokens = tokens.shard(device, 0)
    if (MP := getenv("MP", 1)) > 1:
      device = tuple(f"{Device.DEFAULT}:{i}" for i in range(MP))
      tokens = tokens.shard(device)
-    if DP == 1 and MP == 1: tokens = tokens.to(None)
    logits:Tensor = model(tokens[:, :-1], start_pos=0, temperature=math.nan)
-    loss = vocab_mask.where(-1e9, logits).sparse_categorical_crossentropy(tokens[:, 1:])
+    loss = logits.sparse_categorical_crossentropy(tokens[:, 1:])
    loss.backward()
    assert all(p.grad is g for p,g in zip(optim.params, grads))
    Tensor.realize(loss, *grads)
@@ -1421,37 +1409,35 @@ def train_llama3():

  @TinyJit
  def optim_step():
-    grad_norm = optim.fstep(grads)
+    optim.step()
    scheduler.step()

    for g in grads:
-      g.assign(g.zeros_like()).realize()
+      g.assign(g.zeros_like())

    lr = optim.lr
    Tensor.realize(lr, *grads)

-    return lr.float().to("CPU"), grad_norm.float().to("CPU")
+    return lr.float().to("CPU")

  @TinyJit
  @Tensor.train(False)
  def eval_step(tokens:Tensor):
+    tokens = tokens.to(None)
    if (DP := getenv("DP", 1)) > 1:
      device = tuple(f"{Device.DEFAULT}:{i}" for i in range(DP))
-      tokens = tokens.to(None).shard(device, 0)
+      tokens = tokens.shard(device, 0)
    if (MP := getenv("MP", 1)) > 1:
      device = tuple(f"{Device.DEFAULT}:{i}" for i in range(MP))
      tokens = tokens.shard(device)
-    if DP == 1 and MP == 1: tokens = tokens.to(None)
    logits:Tensor = model(tokens[:, :-1], start_pos=0, temperature=math.nan)
-    loss = vocab_mask.where(-1e9, logits).sparse_categorical_crossentropy(tokens[:, 1:])
+    loss = logits.sparse_categorical_crossentropy(tokens[:, 1:])
    return loss.flatten().float().to("CPU")

  # ** data iters **
  def fake_data(bs, samples):
-    import numpy as np
    for _ in range(samples // bs):
-      fake_data_np = np.random.randint(0, model_params["vocab_size"], size=(bs, SEQLEN + 1), dtype=np.int32)
-      yield Tensor(fake_data_np, device="NPY")
+      yield Tensor.randint(bs, SEQLEN + 1, low=0, high=model_params["vocab_size"], dtype=dtypes.int32, device=Device.DEFAULT)

  def get_train_iter():
    if getenv("FAKEDATA", 0):
@@ -1483,29 +1469,29 @@ def train_llama3():
      st = time.perf_counter()

      stopped = False
-      losses, data_time, dev_time = [], 0, 0
-      for _ in range(grad_acc if i >= 3 else 1):
+      for _ in range(grad_acc):
        ist = time.perf_counter()
        try: tokens = next(train_iter)
        except StopIteration:
          stopped = True
          break
-        mst = time.perf_counter()
-        data_time += mst - ist
-        losses.append(minibatch(tokens).item())
-        dev_time += time.perf_counter() - mst
+        dt = time.perf_counter()
+        loss = minibatch(tokens)
      if stopped: break

      gt = time.perf_counter()
-      ret = optim_step()
-      lr, grad_norm = ret[0].item(), ret[1].item()
-      et = time.perf_counter()
+      lr = optim_step()
+      ot = time.perf_counter()

-      loss = sum(losses) / len(losses)
-      optim_time = et - gt
-      dev_time += optim_time
+      loss = loss.float().item()
+      lr = lr.item()
+
+      et = time.perf_counter()
      step_time = et - st
      gbs_time = gt - st
+      optim_time = ot - gt
+      data_time = dt - ist
+      dev_time = step_time - data_time * grad_acc
      if BENCHMARK: step_times.append(step_time)

      i += 1
@@ -1516,14 +1502,11 @@ def train_llama3():
      mfu = ((6 * num_params * SEQLEN * GBS) / (dev_time * max(getenv("DP", 1), getenv("MP", 1)) * 2.3e15)) * 100
      tqdm.write(
          f"{i:5} {step_time:.3f} s step, {gbs_time:.3f} s gbs, {optim_time:.3f} s optim, {data_time:.3f} s data, {loss:.4f} loss, " \
-          f"{lr:.12f} LR, {grad_norm:.6f} grad_norm, {mem_gb:.2f} GB used, {gflops:9.2f} GFLOPS, {mfu:5.2f}% MFU")
-      if DEBUG >= 1: tqdm.write("  mem per device: " + ', '.join(f"{dev}: {mem/1e9:.2f} GB" for dev, mem in sorted(GlobalCounters.mem_used_per_device.items())))
+          f"{lr:.12f} LR, {mem_gb:.2f} GB used, {gflops:9.2f} GFLOPS, {mfu:5.2f}% MFU")

      if WANDB:
        wandb.log({
-          "train/loss": loss,
-          "train/lr": lr,
-          "train/grad_norm": grad_norm,
+          "lr": lr, "train/loss": loss,
          "train/step_time": step_time,
          "train/gbs_time": gbs_time,
          "train/optim_time": optim_time,
@@ -1560,7 +1543,7 @@ def train_llama3():
      # run eval
      eval_losses = []
      eval_iter = get_eval_iter()
-      tqdm.write(f"evaluating {EVAL_SAMPLES//EVAL_BS} batches of {EVAL_BS} sequences")
+      tqdm.write(f"evaluating {5760//EVAL_BS} batches of {EVAL_BS} sequences")

      for j,tokens in tqdm(enumerate(eval_iter), total=EVAL_SAMPLES//EVAL_BS):
        eval_losses += eval_step(tokens).tolist()
@@ -7,49 +7,39 @@ class GradAccClipAdamW(Optimizer):
  def __init__(self, params:list[Tensor], lr=0.001, b1=0.9, b2=0.999, eps=1e-6, weight_decay=0.0, grad_acc=1, clip_norm=1.0, device=None, fused=FUSE_OPTIM):
    super().__init__(params, lr, device, fused)
    self.b1, self.b2, self.eps, self.wd = b1, b2, eps, weight_decay
-    self.b1_t, self.b2_t = (Tensor.ones((1,), dtype=dtypes.float32, device=self.device, requires_grad=False) for _ in [b1, b2])
+    self.b1_t, self.b2_t = (Tensor.ones((1,), dtype=dtypes.float32, device=self.device, requires_grad=False).contiguous() for _ in [b1, b2])
    self.m = self._new_optim_param()
    self.v = self._new_optim_param()
    self.grad_acc, self.clip_norm = grad_acc, clip_norm

-  def fstep(self, grads:list[Tensor]):
-    if self.fused:
-      out, extra = self._step([], grads)
-      updates = [out[0][self.pos_params[i]:self.pos_params[i+1]].reshape(tt.shape) for i, tt in enumerate(self.params)]
-    else:
-      updates, extra = self._step([], grads)
-    for i, tt in enumerate(self.params): tt.assign(self._apply_update(tt, updates[i]))
-    to_realize = extra+self.params+self.buffers
-
-    Tensor.realize(*to_realize)
-    return extra[-1]
-
  def _step(self, params:list[Tensor], grads:list[Tensor]) -> tuple[list[Tensor], list[Tensor]]:
    for i in range(len(grads)):
-      if grads[i].device != self.m[i].device: grads[i].assign(grads[i].to(self.m[i].device))
+      if grads[i].device != self.m[i].device: grads[i] = grads[i].to(self.m[i].device)

    if self.fused:
-      grads[0].assign(grads[0] / self.grad_acc)
+      grads[0] = grads[0] / self.grad_acc
      total_norm = grads[0].float().square().sum().sqrt()
-      grads[0].assign((grads[0] * (self.clip_norm / (total_norm + 1e-6)).clamp(max_=1.0)).cast(grads[0].dtype))
+      grads[0] = (grads[0] * (self.clip_norm / (total_norm + 1e-6)).clamp(max_=1.0)).cast(grads[0].dtype)
    else:
+      total_norm = Tensor.zeros((), dtype=dtypes.float32, device=self.device)
+      for g in grads:
+        total_norm += g.float().square().sum()
+      total_norm = total_norm.sqrt()
      for i in range(len(grads)):
-        grads[i].assign(grads[i] / self.grad_acc).realize()
-      total_norm = Tensor.stack(*[g.float().square().sum() for g in grads]).sum().sqrt().contiguous().realize()
-      for i in range(len(grads)):
-        grads[i].assign((grads[i] * (self.clip_norm / (total_norm + 1e-6)).clamp(max_=1.0)).cast(grads[i].dtype)).realize()
+        grads[i] = grads[i] / self.grad_acc
+        grads[i] = (grads[i] * (self.clip_norm / (total_norm + 1e-6)).clamp(max_=1.0)).cast(grads[i].dtype)

    ret = []
    self.b1_t *= self.b1
    self.b2_t *= self.b2
-    for i, g in enumerate(grads):
+    for i, (t, g) in enumerate(zip(params, grads)):
      self.m[i].assign((self.b1 * self.m[i] + (1.0 - self.b1) * g).cast(self.m[i].dtype))
      self.v[i].assign((self.b2 * self.v[i] + (1.0 - self.b2) * (g * g)).cast(self.v[i].dtype))
      m_hat = self.m[i] / (1.0 - self.b1_t)
      v_hat = self.v[i] / (1.0 - self.b2_t)
      up = m_hat / (v_hat.sqrt() + self.eps)
-      ret.append((self.lr * up).cast(g.dtype))
-    return ret, [self.b1_t, self.b2_t] + self.m + self.v + [total_norm]
+      ret.append((self.lr * up).cast(t.dtype))
+    return ret, [self.b1_t, self.b2_t] + self.m + self.v

  def _apply_update(self, t:Tensor, up:Tensor) -> Tensor:
    up = up.shard_like(t) + self.lr.to(t.device) * self.wd * t.detach()
@@ -5,7 +5,6 @@ export DEV=${DEV:-AMD}
 export EMULATE="AMD_CDNA4"
 export CHECK_OOB=0
 export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
-export DEVICE_IN_FUNCTION_BUG=1

 export DEBUG=${DEBUG:-0}
 export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
@@ -3,4 +3,4 @@ export BENCHMARK=5
 export EVAL_BS=0
 export VIZ=${VIZ:--1}
 examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/dev_run.sh
-extra/viz/cli.py --profile --device "AMD" --top 20
+PYTHONPATH="." extra/viz/cli.py --profile --device "AMD" --top 20
@@ -1,16 +0,0 @@
-import sys, pickle
-from extra.bench_log import WallTimeEvent, BenchEvent
-from tinygrad.helpers import getenv
-
-PKL = sys.argv[1] if len(sys.argv) > 1 else "/tmp/openpilot.pkl"
-
-load_times = []
-
-for _ in range(10):
-  with WallTimeEvent(BenchEvent.STEP) as wte: pickle.load(open(PKL, 'rb'))
-  load_times.append(wte.time)
-  print(f"pickle load: {wte.time:6.2f} s")
-
-if (assert_time:=getenv("ASSERT_MIN_LOAD_TIME")):
-  min_time = min(load_times)
-  assert min_time < assert_time, f"Speed regression, expected min load time of < {assert_time} s but took: {min_time} s"
@@ -34,8 +34,7 @@ class WallTimeEvent:
    self.start = time.monotonic()
    return self
  def __exit__(self, *_):
-    self.time = time.monotonic() - self.start
-    _events[self.event]["wall"].append(self.time)
+    _events[self.event]["wall"].append(time.monotonic() - self.start)
    return False

 class KernelTimeEvent:
@@ -3,7 +3,7 @@ from tinygrad import Tensor, Device, dtypes
 from tinygrad.dtype import AddrSpace
 from tinygrad.uop.ops import UOp, Ops, KernelInfo, AxisType
 from tinygrad.renderer import Estimates
-from tinygrad.helpers import getenv, all_same, DEBUG
+from tinygrad.helpers import getenv, all_same, dedup
 from extra.gemm.asm.cdna.asm import build_kernel, TILE_M, TILE_N, TILE_K, NUM_WG

 # ** CDNA4 assembly gemm
@@ -11,12 +11,12 @@ from extra.gemm.asm.cdna.asm import build_kernel, TILE_M, TILE_N, TILE_K, NUM_WG
 WORKGROUP_SIZE = 256

@functools.cache
-def custom_asm_gemm(C:UOp, A:UOp, B:UOp, dname:str) -> UOp:
+def custom_asm_gemm(C:UOp, A:UOp, B:UOp, dname:str, arch:str, wg:int) -> UOp:
  batch, M, K = A.shape
  K2, N = B.shape[(1 if B.ndim == 3 else 0):]
  assert K == K2
  lidx = UOp.special(WORKGROUP_SIZE, "lidx0")
-  gidx = UOp.special(NUM_WG, "gidx0")
+  gidx = UOp.special(wg, "gidx0")
  insts = build_kernel(batch, M, N, K, A.dtype.base)
  lds = UOp(Ops.DEFINE_LOCAL, dtypes.uint8.ptr(size=133_120, addrspace=AddrSpace.LOCAL), (), 'lds')
  sink = UOp.sink(C.base, A.base, B.base, lds, lidx, gidx,
@@ -26,25 +26,17 @@ def custom_asm_gemm(C:UOp, A:UOp, B:UOp, dname:str) -> UOp:

 counters = {"used":0, "todos":[]}
 def todo(msg:str) -> bool: counters["todos"].append(msg); return False
-def _asm_gemm_report():
-  print(f'asm_gemm: {counters["used"]} used, {len(counters["todos"])} not used')
-  if DEBUG >= 2 and counters["todos"]:
-    from collections import Counter
-    for msg, cnt in Counter(counters["todos"]).most_common(): print(f'  {cnt:3d}x {msg}')
-atexit.register(_asm_gemm_report)
+atexit.register(lambda: print(f'asm_gemm: {counters["used"]} used, {len(counters["todos"])} not used'))

 def can_use_asm_gemm(a:Tensor, b:Tensor) -> bool:
  if a.dtype != b.dtype: return todo(f"dtypes must match {a.dtype} != {b.dtype}")
  if a.dtype not in {dtypes.bfloat16, dtypes.float16}: return todo(f"only bfloat16/float16, got {a.dtype}")
  batch, M, K = (1, *a.shape) if a.ndim == 2 else a.shape
  N = b.shape[1]
+  # only sharding on the batch or K is tested, others might work too
  if isinstance(a.device, tuple):
-    if a.ndim == 2 and a.uop.axis == 0 and b.uop.axis is None: M //= len(a.device)
-    elif a.ndim == 2 and a.uop.axis == 1 and b.uop.axis == 0: K //= len(a.device)
-    elif a.ndim == 2 and a.uop.axis is None and b.uop.axis == 1: N //= len(a.device)
+    if a.ndim == 2 and a.uop.axis == 1 and b.uop.axis == 0: K //= len(a.device)
    elif a.ndim == 3 and a.uop.axis == 0 and b.uop.axis is None: batch //= len(a.device)
-    elif a.ndim == 3 and a.uop.axis is None and b.uop.axis == 1: N //= len(a.device)
-    elif a.ndim == 3 and a.uop.axis == 2 and b.uop.axis == 0: K //= len(a.device)
    else: return todo(f"sharding mismatch a.ndim={a.ndim} a.uop.axis={a.uop.axis} b.uop.axis={b.uop.axis}")
    dname = a.device[0]
  else: dname = a.device
@@ -86,10 +78,6 @@ def custom_gemm_bw(gradient:UOp, kernel:UOp):
 def asm_gemm(a:Tensor, b:Tensor) -> 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
-  if unfold_batch:
-    orig_batch = a.shape[0]
-    a = a.reshape(a.shape[0]*a.shape[1], a.shape[2])
  squeeze = a.ndim == 2
  if squeeze: a = a.unsqueeze(0)

@@ -97,26 +85,17 @@ def asm_gemm(a:Tensor, b:Tensor) -> Tensor:
  N = b.shape[1]
  is_multi = isinstance(a.device, tuple)
  if (k_sharded:=is_multi and a.uop.axis == 2): K //= len(a.device)
-  if (m_sharded:=is_multi and a.uop.axis == 1): M //= len(a.device)
-  n_sharded = is_multi and b.uop.axis == 1

  if is_multi:
-    if n_sharded:
-      out = Tensor(Tensor.empty(batch, M, N//len(a.device), dtype=a.dtype, device=a.device).uop.multi(2), device=a.device)
-    elif m_sharded:
-      out = Tensor(Tensor.empty(batch, M, N, dtype=a.dtype, device=a.device).uop.multi(1), device=a.device)
-    else:
-      out = Tensor(Tensor.empty(batch//len(a.device) if a.uop.axis==0 else batch, M, N, dtype=a.dtype, device=a.device).uop.multi(0), device=a.device)
+    out = Tensor(Tensor.empty(batch//len(a.device) if a.uop.axis==0 else batch, M, N, dtype=a.dtype, device=a.device).uop.multi(0), device=a.device)
  else:
    out = Tensor.empty(batch, M, N, dtype=a.dtype, device=a.device)

  renderer = Device[a.device[0] if is_multi else a.device].renderer
  dname, arch = renderer.device, getattr(renderer, "arch", "")
  if arch.startswith("gfx950") and getenv("USE_ASM", 1):
-    out = Tensor.custom_kernel(out, a, b, fxn=functools.partial(custom_asm_gemm, dname=dname), grad_fxn=custom_gemm_bw)[0]
+    out = Tensor.custom_kernel(out, a, b, fxn=functools.partial(custom_asm_gemm, dname=dname, wg=NUM_WG, arch=arch), grad_fxn=custom_gemm_bw)[0]
  else:
    out = Tensor.custom_kernel(out, a, b, fxn=custom_uop_gemm, grad_fxn=custom_gemm_bw)[0]
  if k_sharded: out = out.sum(0)
-  out = out.squeeze(0) if squeeze else out
-  if unfold_batch: out = out.reshape(orig_batch, -1, out.shape[-1])
-  return out
+  return out.squeeze(0) if squeeze else out
@@ -56,10 +56,7 @@ class Attention:
  def __call__(self, x:Tensor, start_pos:Union[Variable,int], freqs_cis:Tensor, mask:Optional[Tensor]=None) -> Tensor:
    if getenv("WQKV"):
      xqkv = self.wqkv(x)
-      xqkv = xqkv.reshape(xqkv.shape[0], xqkv.shape[1], self.n_kv_heads, self.n_rep + 2, self.head_dim)
-      xq = xqkv[:, :, :, :self.n_rep].reshape(xqkv.shape[0], xqkv.shape[1], -1)
-      xk = xqkv[:, :, :, self.n_rep:self.n_rep+1].reshape(xqkv.shape[0], xqkv.shape[1], -1)
-      xv = xqkv[:, :, :, self.n_rep+1:self.n_rep+2].reshape(xqkv.shape[0], xqkv.shape[1], -1)
+      xq, xk, xv = xqkv.split([self.n_heads * self.head_dim, self.n_kv_heads * self.head_dim, self.n_kv_heads * self.head_dim], dim=2)
    else:
      xq, xk, xv = self.wq(x), self.wk(x.contiguous_backward()), self.wv(x)

@@ -2,13 +2,30 @@

 ## Getting SQ Thread Trace

-`VIZ=2` to enable SQTT profiling.
-
-`SQTT_ITRACE_SE_MASK=X` to select shader engines for instruction tracing, -1 = all, 0 = disabled, >0 = SE bitmask, default 0b11.
+SQTT is implemented on top of normal tinygrad profiling, `VIZ=1 SQTT=1` to get profile pickle with sqtt data embedded in it.

 `SQTT_BUFFER_SIZE=X` to change size of SQTT buffer (per shader engine, 6 SEs on 7900xtx) in megabytes, default 256.

-## Viewing the traces
+`SQTT_ITRACE_SE_MASK=X` to select for which shader engines instruction tracing will be enabled, -1 is all, 0 is none (instruction tracing disabled), >0 is
+bitfield/mask for SEs to enable instruction tracing on. Masking shader engines will give smaller file sizes at a cost of less hits and kernels that
+don't have any wavefront on first simd of shader engine with instruction tracing enabled will not have instruction timings.
+The default is 2 (second shader engine only), only one for file size reasons, second instead of first because dispatch starts from it so there is
+greater chance that kernels with small global size will have instruction tracing data.
+ 
+Note that instruction tracing might not be available for kernels with small global dims, this is not a bug, but it can be improved with various hacks
+to the point where it can reliably trace a kernel consisting of a single wavefront (am only, not quite reliable under amdgpu due to waves sometimes
+being dispatched starting from different simds). More info in comments in ops_amd.py

- Web UI: `tinygrad/viz/serve.py`
- Command line: `python -m tinygrad.renderer.amd.sqtt`
+## Converting pickled profile with SQTT data into RGP file
+
+```bash
+extra/sqtt/rgptool.py create "/tmp/profile.pkl.$USER" -o /tmp/gpu0.rgp
+```
+
+Then load gpu0.rgp into Radeon GPU Profiler. It works just fine both in wine (macos, native version available for linux) and via ssh X forwarding
+
+If multiple gpus are used you can select which one to export with `-d` like this:
+
+```bash
+extra/sqtt/rgptool.py create "/tmp/profile.pkl.$USER" -d 'AMD:5' -o /tmp/gpu5.rgp
+```
@@ -0,0 +1,152 @@
+import os
+os.environ["PYTHONPATH"] = "."
+os.environ["SQTT"] = "1"
+if "DEV" not in os.environ: os.environ["DEV"] = "AMD"
+os.environ["PROFILE"] = "1"
+os.environ["AMD_LLVM"] = "0"
+
+from dataclasses import replace
+import atexit, contextlib
+from tinygrad import Tensor
+from tinygrad.helpers import system, OSX
+from tinygrad.runtime.ops_amd import AMDProgram
+from extra.sqtt.roc import decode, WaveExec, ProfileSQTTEvent
+from tinygrad.device import Device
+
+from extra.sqtt.attempt_sqtt_parse import parse_sqtt_print_packets
+
+dev = Device["AMD"]
+
+@contextlib.contextmanager
+def save_sqtt():
+  # clear the old traces
+  dev.profile_events.clear()
+  sqtt:dict[str, list[WaveExec]] = {}
+  yield sqtt
+  events = dev.profile_events
+
+  #rctx = decode(events)
+  #assert len(rctx.inst_execs) > 0, "empty sqtt output"
+  #sqtt.update(rctx.inst_execs)
+
+  for e in events:
+    if isinstance(e, ProfileSQTTEvent):
+      print(replace(e, blob=b''))
+      if e.se == 0:
+        parse_sqtt_print_packets(e.blob)
+
+template = """.text
+.globl matmul
+.p2align 8
+.type matmul,@function
+matmul:
+  INSTRUCTION
+
+.rodata
+.p2align 6
+.amdhsa_kernel matmul
+  .amdhsa_kernarg_size 8
+  .amdhsa_user_sgpr_kernarg_segment_ptr 1
+  .amdhsa_next_free_vgpr .amdgcn.next_free_vgpr
+  .amdhsa_next_free_sgpr .amdgcn.next_free_sgpr
+  .amdhsa_wavefront_size32 1
+.end_amdhsa_kernel
+
+.amdgpu_metadata
+---
+amdhsa.version:
+  - 1
+  - 0
+amdhsa.kernels:
+  - .name: matmul
+    .symbol: matmul.kd
+    .group_segment_fixed_size: 0
+    .private_segment_fixed_size: 0
+    .wavefront_size: 32
+    .sgpr_count: 8
+    .vgpr_count: 8
+    .max_flat_workgroup_size: 1024
+    .kernarg_segment_align: 8
+    .kernarg_segment_size: 8
+    .args:
+      - .address_space:  global
+        .name:           a
+        .offset:         0
+        .size:           8
+        .type_name:      'float*'
+        .value_kind:     global_buffer
+...
+.end_amdgpu_metadata
+"""
+
+def run_asm(src, num_workgroups=1, num_waves=1):
+  WAVE_SIZE = 32
+  t = Tensor.empty(0x1000).realize()
+  buf = t.uop.buffer.ensure_allocated()
+  lib = dev.compiler.compile(template.replace("INSTRUCTION", '\n'.join(src)))
+  dev.compiler.disassemble(lib)
+  fxn = AMDProgram(dev, "matmul", lib)
+  fxn(buf._buf, global_size=(num_workgroups,1,1), local_size=(WAVE_SIZE*num_waves,1,1), wait=True)
+
+if __name__ == "__main__":
+  with save_sqtt() as sqtt:
+    run_asm([
+      "s_nop 100",
+      "s_nop 100",
+      "s_load_b64 s[0:1], s[0:1], null",
+      "s_waitcnt lgkmcnt(0)",
+      "s_nop 100",
+      "s_nop 100",
+      "s_add_i32 s2, s2, 10",
+      "s_add_i32 s2, s2, 10",
+      "s_nop 100",
+      "s_nop 100",
+      "v_mov_b32_e32 v0, 0",
+      "v_mov_b32_e32 v0, 0",
+      "s_nop 100",
+      "s_nop 100",
+      "v_dual_fmac_f32 v2, v48, v24 :: v_dual_fmac_f32 v9, v37, v51",
+      "v_dual_fmac_f32 v2, v48, v24 :: v_dual_fmac_f32 v9, v37, v51",
+      "s_nop 100",
+      "s_nop 100",
+      "global_load_b128 v[2:5], v0, s[0:1]",
+      "global_load_b128 v[2:5], v0, s[0:1]",
+      "s_nop 100",
+      "s_nop 100",
+      "s_sendmsg sendmsg(MSG_DEALLOC_VGPRS)",
+      "s_endpgm",
+    ], num_workgroups=1, num_waves=1)
+  exit(0)
+
+  with save_sqtt() as sqtt:
+    #(Tensor.empty(16,16) @ Tensor.empty(16,16)).elu().realize()
+    #Tensor.empty(1, 64).sum(axis=1).realize()
+    Tensor.empty(1).log2().realize()
+  exit(0)
+
+  with save_sqtt() as sqtt:
+    # what's in v0?
+    run_asm([
+      "v_mov_b32_e32 v0, 0",
+      "v_mov_b32_e32 v1, 0",
+      "s_clause 0x1",
+      "s_load_b64 s[0:1], s[0:1], null",
+      "s_waitcnt lgkmcnt(0)",
+    ]+[
+      "global_load_b32 v1, v0, s[0:1]",
+    ]*10+[
+      "global_load_b32 v10, v1, s[0:1]",
+      "s_waitcnt vmcnt(0)",
+
+      #"v_rcp_f32 v1, v0"
+      #"v_add_f32_e32 v1 v0 v0",
+      #"v_add_f32_e32 v5 v4 v4",
+      #"v_add_f32_e32 v7 v6 v6",
+      #"v_add_f32_e32 v1 v0 v0",
+      #"v_add_f32_e32 v2 v1 v1",
+      #"s_nop 1"
+    ]*5+[
+      "v_add_f32_e32 v3 v2 v2",
+    ]*5+[
+      "v_mul_f32_e32 v3 v2 v2",
+    ]*7)
@@ -0,0 +1,548 @@
+import pickle, sys
+from tinygrad.helpers import getenv, Timing, colored
+from extra.sqtt.roc import decode, ProfileSQTTEvent
+
+# do these enums match fields in the packets?
+#from tinygrad.runtime.support.amd import import_soc
+#soc = import_soc([11])
+#perf_sel = {getattr(soc, k):k for k in dir(soc) if k.startswith("SQ_PERF_")}
+
+# Instruction packets (one per ISA op)
+# NOTE: these are bad guesses and may be wrong! feel free to update if you know better
+# some names were taken from SQ_TT_TOKEN_MASK_TOKEN_EXCLUDE_SHIFT
+
+# we see 18 opcodes
+# opcodes(18):  1  2  3  4  5  6  8  9  F 10 11 12 14 15 16 17 18 19
+# if you exclude everything, you are left with 6
+# opcodes( 6): 10 11 14 15 16 17
+# sometimes we see a lot of B, but not repeatable
+
+# not seen
+# 7 A C
+
+# NOTE: INST runs before EXEC
+
+OPCODE_COLORS = {
+  # dispatches are BLACK
+  0x1: "BLACK",
+  0x18: "BLACK",
+
+  # execs are yellow
+  0x2: "yellow",
+  0x3: "yellow",
+  0x4: "YELLOW",
+  0x5: "YELLOW",
+
+  # waves are blue
+  0x8: "blue",
+  0x9: "blue",
+  0x6: "cyan",
+  0xb: "cyan",
+}
+
+OPCODE_NAMES = {
+  # gated by SQ_TT_TOKEN_EXCLUDE_VALUINST_SHIFT (but others must be enabled for it to show)
+  0x01: "VALUINST",
+  # gated by SQ_TT_TOKEN_EXCLUDE_VMEMEXEC_SHIFT
+  0x02: "VMEMEXEC",
+  # gated by SQ_TT_TOKEN_EXCLUDE_ALUEXEC_SHIFT
+  0x03: "ALUEXEC",
+  # gated by SQ_TT_TOKEN_EXCLUDE_IMMEDIATE_SHIFT
+  0x04: "IMMEDIATE",
+  0x05: "IMMEDIATE_MASK",
+
+  # gated by SQ_TT_TOKEN_EXCLUDE_WAVERDY_SHIFT
+  0x06: "WAVERDY",
+  # gated by SQ_TT_TOKEN_EXCLUDE_WAVESTARTEND_SHIFT
+  0x08: "WAVEEND",
+  0x09: "WAVESTART",
+  # gated by SQ_TT_TOKEN_EXCLUDE_WAVEALLOC_SHIFT
+  0x0B: "WAVEALLOC",  # FFF00
+
+  # gated by NOT SQ_TT_TOKEN_EXCLUDE_PERF_SHIFT
+  0x0D: "PERF",
+  # gated by SQ_TT_TOKEN_EXCLUDE_EVENT_SHIFT
+  0x12: "EVENT",
+  0x13: "EVENT_BIG",  # FFFFF800
+  # some gated by SQ_TT_TOKEN_EXCLUDE_REG_SHIFT, some always there. something is broken with the timing on this
+  0x14: "REG",
+  # gated by SQ_TT_TOKEN_EXCLUDE_INST_SHIFT
+  0x18: "INST",
+  # gated by SQ_TT_TOKEN_EXCLUDE_UTILCTR_SHIFT
+  0x19: "UTILCTR",
+
+  # this is the first (8 byte) packet in the bitstream
+  0x17: "LAYOUT_HEADER",       # layout/mode/group + selectors A/B (reversed)
+
+  # pure time (no extra bits)
+  0x0F: "TS_DELTA_SHORT",
+  0x10: "NOP",
+  0x11: "TS_WAVE_STATE",     # almost pure time, has a small flag
+
+  # not a good name, but seen and understood mostly
+  0x15: "SNAPSHOT",          # small delta + 50-ish bits of snapshot
+  0x16: "TS_DELTA_OR_MARK",  # 36-bit long delta or 36-bit marker
+
+  # packets we haven't seen / rarely see 0x0b
+  0x07: "TS_DELTA_S8_W3_7",         # shift=8,  width=3  (small delta)
+  0x0A: "TS_DELTA_S5_W2_A",         # shift=5,  width=2
+  0x0C: "TS_DELTA_S5_W3_B",         # shift=5,  width=3 (different consumer)
+}
+
+#  SALU    =  0x0 / s_mov_b32
+#  SMEM    =  0x1 / s_load_b*
+#  JUMP    =  0x3 / s_cbranch_scc0
+#  NEXT    =  0x4 / s_cbranch_execz
+#  MESSAGE =  0x9 / s_sendmsg
+#  VALU    =  0xb / v_(exp,log)_f32_e32
+#  VALU    =  0xd / v_lshlrev_b64
+#  VALU    =  0xe / v_mad_u64_u32
+#  VMEM    = 0x21 / global_load_b32
+#  VMEM    = 0x22 / global_load_b32
+#  VMEM    = 0x24 / global_store_b32
+#  VMEM    = 0x25 / global_store_b64
+#  VMEM    = 0x27 / global_store
+#  VMEM    = 0x28 / global_store_b64
+#  LDS     = 0x29 / ds_load_b128
+#  LDS     = 0x2b / ds_store_b32
+#  LDS     = 0x2e / ds_store_b128
+#  ????    = 0x5a / hidden global_load  instruction
+#  ????    = 0x5b / hidden global_load  instruction
+#  ????    = 0x5c / hidden global_store instruction
+#  VALU    = 0x73 / v_cmpx_eq_u32_e32 (not normal VALUINST)
+OPNAME = {
+  0x0: "SALU",
+  0x1: "SMEM",
+  0x3: "JUMP",
+  0x4: "NEXT",
+  0x9: "MESSAGE",
+  0xb: "VALU",
+  0xd: "VALU",
+  0xe: "VALU",
+  0x21: "VMEM_LOAD",
+  0x22: "VMEM_LOAD",
+  0x24: "VMEM_STORE",
+  0x25: "VMEM_STORE",
+  0x26: "VMEM_STORE",
+  0x27: "VMEM_STORE",
+  0x28: "VMEM_STORE",
+  0x29: "LDS_LOAD",
+  0x2b: "LDS_STORE",
+  0x2e: "LDS_STORE",
+  0x50: "__SIMD_LDS_LOAD",
+  0x51: "__SIMD_LDS_LOAD",
+  0x54: "__SIMD_LDS_STORE",
+  0x5a: "__SIMD_VMEM_LOAD",
+  0x5b: "__SIMD_VMEM_LOAD",
+  0x5c: "__SIMD_VMEM_STORE",
+  0x5d: "__SIMD_VMEM_STORE",
+  0x5e: "__SIMD_VMEM_STORE",
+  0x5f: "__SIMD_VMEM_STORE",
+  0x72: "SALU_OR",
+  0x73: "VALU_CMPX",
+}
+
+ALUSRC = {
+  1: "SALU",
+  2: "VALU",
+  3: "VALU_SALU",
+}
+
+MEMSRC = {
+  0: "LDS",
+  1: "__LDS",
+  2: "VMEM",
+  3: "__VMEM",
+}
+
+
+# these tables are from rocprof trace decoder
+# rocprof_trace_decoder_parse_data-0x11c6a0
+# parse_sqtt_180 = b *rocprof_trace_decoder_parse_data-0x11c6a0+0x110040
+
+# ---------- 1. local_138: 256-byte state->opcode table ----------
+
+STATE_TO_OPCODE: bytes = bytes([
+  0x10, 0x16, 0x18, 0x01, 0x05, 0x0b, 0x0c, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x09, 0x04, 0x03, 0x02,
+  0x10, 0x17, 0x18, 0x01, 0x06, 0x08, 0x0d, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x0a, 0x04, 0x03, 0x02,
+  0x10, 0x07, 0x18, 0x01, 0x05, 0x0b, 0x0c, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x09, 0x04, 0x03, 0x02,
+  0x10, 0x19, 0x18, 0x01, 0x06, 0x08, 0x0d, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x0a, 0x04, 0x03, 0x02,
+  0x10, 0x00, 0x18, 0x01, 0x05, 0x0b, 0x0c, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x09, 0x04, 0x03, 0x02,
+  0x10, 0x11, 0x18, 0x01, 0x06, 0x08, 0x0d, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x0a, 0x04, 0x03, 0x02,
+  0x10, 0x12, 0x18, 0x01, 0x05, 0x0b, 0x0c, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x09, 0x04, 0x03, 0x02,
+  0x10, 0x15, 0x18, 0x01, 0x06, 0x08, 0x0d, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x0a, 0x04, 0x03, 0x02,
+  0x10, 0x16, 0x18, 0x01, 0x05, 0x0b, 0x0c, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x09, 0x04, 0x03, 0x02,
+  0x10, 0x17, 0x18, 0x01, 0x06, 0x08, 0x0d, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x0a, 0x04, 0x03, 0x02,
+  0x10, 0x07, 0x18, 0x01, 0x05, 0x0b, 0x0c, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x09, 0x04, 0x03, 0x02,
+  0x10, 0x19, 0x18, 0x01, 0x06, 0x08, 0x0d, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x0a, 0x04, 0x03, 0x02,
+  0x10, 0x00, 0x18, 0x01, 0x05, 0x0b, 0x0c, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x09, 0x04, 0x03, 0x02,
+  0x10, 0x11, 0x18, 0x01, 0x06, 0x08, 0x0d, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x0a, 0x04, 0x03, 0x02,
+  0x10, 0x13, 0x18, 0x01, 0x05, 0x0b, 0x0c, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x09, 0x04, 0x03, 0x02,
+  0x10, 0x15, 0x18, 0x01, 0x06, 0x08, 0x0d, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x0a, 0x04, 0x03, 0x02,
+])
+
+# opcode mask (the bits used to determine the opcode, worked out by looking at the repeats in STATE_TO_OPCODE)
+
+opcode_mask = {
+  0x10: 0b1111,
+
+  0x16: 0b1111111,
+  0x17: 0b1111111,
+  0x07: 0b1111111,
+  0x19: 0b1111111,
+  0x11: 0b1111111,
+  0x12: 0b11111111,
+  0x13: 0b11111111,
+  0x15: 0b1111111,
+
+  0x18: 0b111,
+  0x1: 0b111,
+
+  0x5: 0b11111,
+  0x6: 0b11111,
+  0xb: 0b11111,
+  0x8: 0b11111,
+  0xc: 0b11111,
+  0xd: 0b11111,
+
+  0xf: 0b1111,
+  0x14: 0b1111,
+
+  0x9: 0b11111,
+  0xa: 0b11111,
+
+  0x4: 0b1111,
+  0x3: 0b1111,
+  0x2: 0b1111,
+}
+
+# ---------- 2. DAT_0012e280: nibble budget per opcode&0x1F ----------
+
+NIBBLE_BUDGET = [
+  0x08, 0x0C, 0x08, 0x08, 0x0C, 0x18, 0x18, 0x40, 0x14, 0x20, 0x30, 0x14, 0x34, 0x1C, 0x30, 0x08,
+  0x04, 0x18, 0x18, 0x20, 0x40, 0x40, 0x30, 0x40, 0x14, 0x30, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+]
+
+# ---------- 3. delta_map from your hash nodes ----------
+
+# opcode -> (shift, width)
+DELTA_MAP_DEFAULT = {
+  0x01: (3,  3),   # shift=3,  end=6
+  0x02: (4,  2),   # shift=4,  end=6
+  0x03: (4,  2),   # shift=4,  end=6
+  0x04: (4,  3),   # shift=4,  end=7
+  0x05: (5,  3),   # shift=5,  end=8
+  0x06: (5,  3),   # shift=5,  end=8
+  0x07: (8,  3),   # shift=8,  end=11
+  0x08: (5,  3),   # shift=5,  end=8
+  0x09: (5,  2),   # shift=5,  end=7
+  0x0A: (5,  2),   # shift=5,  end=7
+  0x0B: (5,  3),   # shift=5,  end=8
+  0x0C: (5,  3),   # shift=5,  end=8
+  0x0D: (5,  3),   # shift=5,  end=8
+  # NOTE: 0x0e can never be decoded, it's not in the STATE_TO_OPCODE table
+  #0x0E: (7,  2),   # shift=7,  end=9
+  0x0F: (4,  4),   # shift=4,  end=8
+  0x10: (0,  0),   # shift=0,  end=0  (no delta)
+  0x11: (7,  9),   # shift=7,  end=16
+  0x12: (8,  3),   # shift=8,  end=11
+  0x13: (8,  3),   # shift=8,  end=11
+  0x14: (4,  3),   # shift=4,  end=7
+  0x15: (7,  3),   # shift=7,  end=10
+  0x16: (12, 36),  # shift=12, end=48 (36-bit field, matches the 0x16 special-case)
+  0x17: (0,  0),   # shift=0,  end=0  (no delta)
+  0x18: (4,  3),   # shift=4,  end=7
+  0x19: (7,  2),   # shift=7,  end=9
+}
+
+# ---------- 4. One-line-per-packet parser ----------
+
+def reg_mask(opcode):
+  nb_bits = NIBBLE_BUDGET[opcode & 0x1F]
+  shift, width = DELTA_MAP_DEFAULT[opcode]
+  delta_mask = ((1 << width) - 1) << shift
+  assert delta_mask & opcode_mask[opcode] == 0, "masks shouldn't overlap"
+  return ((1 << nb_bits) - 1) & ~(delta_mask | opcode_mask[opcode])
+
+def decode_packet_fields(opcode: int, reg: int) -> str:
+  """
+  Decode packet payloads conservatively, using:
+    - NIBBLE_BUDGET[opcode & 0x1F] to mask reg down to true width.
+    - DELTA_MAP_DEFAULT[opcode] to expose the "primary" field (often delta).
+    - Per-opcode layouts derived from rocprof's decompiled consumers.
+  """
+  # --- 0. Restrict to real packet bits not used in delta ---------------------------------
+  pkt = reg & reg_mask(opcode)
+  fields: list[str] = []
+
+  match opcode:
+    case 0x01: # VALUINST
+      # 6 bit field
+      flag = (pkt >> 6) & 1
+      wave = pkt >> 7
+      fields.append(f"wave={wave:x}")
+      if flag: fields.append("flag")
+    case 0x02: # VMEMEXEC
+      # 2 bit field (pipe is a guess)
+      src = pkt>>6
+      fields.append(f"src={src} [{MEMSRC.get(src, '')}]")
+    case 0x03: # ALUEXEC
+      # 2 bit field
+      src = pkt>>6
+      fields.append(f"src={src} [{ALUSRC.get(src, '')}]")
+    case 0x04: # IMMEDIATE_4
+      # 5 bit field (actually 4)
+      wave = pkt >> 7
+      fields.append(f"wave={wave:x}")
+    case 0x05: # IMMEDIATE_5
+      # 16 bit field
+      # 1 bit per wave
+      fields.append(f"mask={pkt>>8:016b}")
+    case 0x6:
+      # wave ready FFFF00
+      # 16 bit field
+      # 1 bit per wave
+      fields.append(f"mask={pkt>>8:016b}")
+    case 0x0d:
+      # 20 bit field
+      fields.append(f"arg = {pkt>>8:X}")
+    case 0x12:
+      fields.append(f"event = {pkt>>11:X}")
+    case 0x15:
+      fields.append(f"snap = {pkt>>10:X}")
+    case 0x19:
+      # wave end
+      fields.append(f"ctr = {pkt>>9:X}")
+    case 0xf:
+      extracted_delta = (reg >> 4) & 0xF
+      fields.append(f"strange_delta=0x{extracted_delta:x}")
+    case 0x11:
+      # DELTA_MAP_DEFAULT: shift=7, width=9 -> small delta.
+      # FF0000 is the mask
+      coarse    = pkt >> 16
+      fields.append(f"coarse=0x{coarse:02x}")
+      # From decomp:
+      #  - when layout<3 and coarse&1, it sets a "has interesting wave" flag
+      #  - when coarse&8, it marks all live waves as "terminated"
+      if coarse & 0x01:
+        fields.append("flag_wave_interest=1")
+      if coarse & 0x08:
+        fields.append("flag_terminate_all=1")
+    case 0x8:
+      # wave end, this is 20 bits (FFF00)
+      flag7   = (pkt >> 8) & 1
+      simd    = (pkt >> 9) & 3
+      cu      = ((pkt >> 11) & 0x7) | (flag7 << 3)
+      wave    = (pkt >> 15) & 0x1f
+      fields.append(f"wave={wave:x}")
+      fields.append(f"simd={simd}")
+      fields.append(f"cu={cu}")
+    case 0x9:
+      # From case 9 (WAVESTART) in multiple consumers:
+      #   flag7  = (w >> 7) & 1        (low bit of uVar41)
+      #   cls2   = (w >> 8) & 3        (class / group)
+      #   slot4  = (w >> 10) & 0xf     (slot / group index)
+      #   idx_lo = (w >> 0xd) & 0x1f   (low index, layout<4 path)
+      #   idx_hi = (w >> 0xf) & 0x1f   (high index, layout>=4 path)
+      #   id7    = (w >> 0x19) & 0x7f  (7-bit id)
+      flag7   = (pkt >> 7) & 1
+      simd    = (pkt >> 8) & 3
+      cu      = ((pkt >> 10) & 0x7) | (flag7 << 3)
+      wave    = (pkt >> 13) & 0x1F
+      id7     = (pkt >> 17)
+      fields.append(f"wave={wave:x}")
+      fields.append(f"simd={simd}")
+      fields.append(f"cu={cu}")
+      fields.append(f"id7=0x{id7:x}")
+    case 0x18:
+      # FFF88 is the mask
+      # From case 0x18:
+      #   low3   = w & 7
+      #   grp3   = (w >> 3) or (w >> 4) & 7   (layout-dependent)
+      #   flags  = bits 6 (B6) and 7 (B7)
+      #   hi8    = (w >> 0xc) & 0xff   (layout 4 path)
+      #   hi7    = (w >> 0xd) & 0x7f   (other layouts)
+      #   idx5   = (w >> 7) or (w >> 8) & 0x1f, used as wave index
+      flag1 = (pkt >> 3) & 1
+      flag2 = (pkt >> 7) & 1
+      wave = (pkt >> 8) & 0x1F
+      op = (pkt >> 13)
+      fields.append(f"wave={wave:x}")
+      fields.append(f"op=0x{op:02x} [{OPNAME.get(op, '')}]")
+      if flag1: fields.append("flag1")
+      if flag2: fields.append("flag2")
+    case 0x14:
+      subop   = (pkt >> 16) & 0xFFFF       # (short)(w >> 0x10)
+      val32   = (pkt >> 32) & 0xFFFFFFFF   # (uint)(w >> 0x20)
+      slot    = (pkt >> 7) & 0x7           # index in local_168[...] tables
+      hi_byte = (pkt >> 8) & 0xFF          # determines config vs marker
+
+      fields.append(f"subop=0x{subop:04x}")
+      fields.append(f"slot={slot}")
+      fields.append(f"val32=0x{val32:08x}")
+
+      if hi_byte & 0x80:
+        # Config flavour: writes config words into per-slot state arrays.
+        fields.append("kind=config")
+        if subop == 0x000C:
+          fields.append("slot=lo")
+        elif subop == 0x000D:
+          fields.append("slot=hi")
+      else:
+        # COR marker: subop 0xC342, payload "COR\0" → start of a COR region.
+        if subop == 0xC342:
+          fields.append("kind=cor_stream")
+          if val32 == 0x434F5200:
+            fields.append("cor_magic='COR\\0'")
+    case 0x16:
+      # Bits:
+      #   bit8  -> 0x100
+      #   bit9  -> 0x200
+      #   bits 12..47 -> 36-bit field used as delta or marker
+      bit8 = bool(pkt & 0x100)
+      bit9 = bool(pkt & 0x200)
+      if not bit9:
+        mode = "delta"
+      elif not bit8:
+        mode = "marker"
+      else:
+        mode = "other"
+      # need to use reg here
+      val36 = (reg >> 12) & ((1 << 36) - 1)
+      fields.append(f"mode={mode}")
+      if mode != "delta":
+        fields.append(f"val36=0x{val36:x}")
+    case 0x17:
+      # From decomp (two sites with identical logic):
+      #   layout = (w >> 7) & 0x3f
+      #   mode   = (w >> 0xd) & 3
+      #   group  = (w >> 0xf) & 7
+      #   sel_a  = (w >> 0x1c) & 0xf
+      #   sel_b  = (w >> 0x21) & 7
+      #   flag4  = (w >> 0x3b) & 1  (only meaningful when layout == 4)
+      layout = (pkt >> 7)  & 0x3F
+      simd   = (pkt >> 13) & 0x3  # you can change this by changing traced simd
+      group  = (pkt >> 15) & 0x7
+      sel_a  = (pkt >> 0x1C) & 0xF
+      sel_b  = (pkt >> 0x21) & 0x7
+      flag4  = (pkt >> 0x3B) & 0x1
+
+      fields.append(f"layout={layout}")
+      fields.append(f"group={group}")
+      fields.append(f"simd={simd}")
+      fields.append(f"sel_a={sel_a}")
+      fields.append(f"sel_b={sel_b}")
+      if layout == 4:
+        fields.append(f"layout4_flag={flag4}")
+    case _:
+      fields.append(f"{pkt:X} & {reg_mask(opcode):X}")
+  return ",".join(fields)
+
+FILTER_LEVEL = getenv("FILTER", 1)
+
+DEFAULT_FILTER: tuple[int, ...] = tuple()
+# NOP + pure time + "sample"
+if FILTER_LEVEL >= 0: DEFAULT_FILTER += (0x10, 0xf, 0x11)
+# reg + event + sample + marker
+# TODO: events are probably good
+if FILTER_LEVEL >= 1: DEFAULT_FILTER += (0x14, 0x12, 0x16)
+# instruction runs + valuinst
+if FILTER_LEVEL >= 2: DEFAULT_FILTER += (0x01, 0x02, 0x03)
+# instructions dispatch (inst, immed)
+if FILTER_LEVEL >= 3: DEFAULT_FILTER += (0x4, 0x5, 0x18)
+# waves
+if FILTER_LEVEL >= 4: DEFAULT_FILTER += (0x6, 0x8, 0x9)
+
+def parse_sqtt_print_packets(data: bytes, filter=DEFAULT_FILTER, verbose=True) -> None:
+  """
+  Minimal debug: print ONE LINE per decoded token (packet).
+
+  Now prints only the actual nibbles that belong to each packet, instead of
+  the full 64-bit shift register.
+  """
+  n = len(data)
+  time = 0
+  last_printed_time = 0
+  reg = 0          # shift register
+  offset = 0       # bit offset, in steps of 4 (one nibble)
+  nib_budget = 0x40
+  flags = 0
+  token_index = 0
+  opcodes_seen = set()
+
+  while (offset >> 3) < n:
+    # 1) Fill register with nibbles according to nib_budget
+    if nib_budget != 0:
+      target = offset + 4 + ((nib_budget - 1) & ~3)
+      while offset != target and (offset >> 3) < n:
+        byte = data[offset >> 3]
+        nib = (byte >> (offset & 4)) & 0xF
+        reg = ((reg >> 4) | (nib << 60)) & ((1 << 64) - 1)
+        offset += 4
+      if offset != target: break  # don't parse past the end
+
+    # 2) Decode token from low 8 bits
+    opcode = STATE_TO_OPCODE[reg & 0xFF]
+    opcodes_seen.add(opcode)
+
+    # 4) Set next nibble budget based on opcode
+    nib_budget = NIBBLE_BUDGET[opcode & 0x1F]
+
+    # 5) Get delta
+    shift, width = DELTA_MAP_DEFAULT[opcode]
+    delta = (reg >> shift) & ((1 << width) - 1)
+
+    # 6) Update time and handle special opcodes 0xF/0x16
+    if opcode == 0x16:
+      two_bits = (reg >> 8) & 0x3
+      if two_bits == 1:
+        flags |= 0x01
+
+      # Common 36-bit field at bits [12..47]
+      if (reg & 0x200) == 0:
+        # delta mode: add 36-bit delta to time
+        pass
+      elif (reg & 0x100) == 0:
+        # marker / other modes: no time advance
+        # real marker: bit9=1, bit8=0, non-zero payload
+        # "other" 0x16 variants, ignored for timing
+        delta = 0
+      else:
+        raise RuntimeError("unknown 0x16 delta")
+    elif opcode == 0x0F:
+      # opcode 0x0F has an offset of 4 to the delta
+      # update: it's actually computed to be 8 to match WAVESTART
+      delta = delta + 8
+
+    # Append extra decoded fields into the note string
+    note = decode_packet_fields(opcode, reg)
+
+    # this delta happens before the instruction
+    time += delta
+    token_index += 1
+
+    if verbose and (filter is None or opcode not in filter):
+      print(f"{time:8d} +{time-last_printed_time:8d} : "+colored(f"{OPCODE_NAMES[opcode]:18s} ", OPCODE_COLORS.get(opcode, "white"))+f"{note}")
+      last_printed_time = time
+
+  # Optional summary at the end
+  print(f"# done: tokens={token_index:_}, final_time={time}, flags=0x{flags:02x}")
+  if verbose:
+    print(f"opcodes({len(opcodes_seen):2d}):",
+          ' '.join([colored(f"{op:2X}", "WHITE" if op in opcodes_seen else "BLACK") for op in sorted(opcode_mask)]))
+
+
+def parse(fn:str):
+  with Timing(f"unpickle {fn}: "): dat = pickle.load(open(fn, "rb"))
+  #if getenv("ROCM", 0):
+  #  with Timing(f"decode {fn}: "): ctx = decode(dat)
+  dat_sqtt = [x for x in dat if isinstance(x, ProfileSQTTEvent)]
+  print(f"got {len(dat_sqtt)} SQTT events in {fn}")
+  return dat_sqtt
+
+if __name__ == "__main__":
+  fn = "extra/sqtt/examples/profile_gemm_run_0.pkl"
+  dat_sqtt = parse(sys.argv[1] if len(sys.argv) > 1 else fn)
+  for i,dat in enumerate(dat_sqtt):
+    with Timing(f"decode pkt {i} with len {len(dat.blob):_}: "):
+      parse_sqtt_print_packets(dat.blob, verbose=getenv("V", 1))
@@ -1,148 +0,0 @@
-#!/usr/bin/env python3
-# Run all ALU and memory instructions in the ISA
-import functools, inspect
-from enum import Enum
-from tinygrad import Tensor, Device, dtypes
-from tinygrad.uop.ops import UOp, Ops, KernelInfo, AddrSpace
-from tinygrad.renderer.amd.dsl import Inst, Reg, OPERANDS, SrcField, VGPRField, SGPRField, SSrcField, SBaseField, AlignedSGPRField, BitField
-from tinygrad.renderer.amd.dsl import FixedBitField, EnumBitField, s, v, NULL, VCC_LO
-from extra.gemm.amd_asm_matmul import Kernel
-
-# skip instructions that mutate wave state (PC, EXEC, allocations, signals)
-SKIP = {"S_SETPC_B64", "S_SWAPPC_B64", "S_RFE_B64", "S_BARRIER_SIGNAL_ISFIRST", "S_GET_BARRIER_STATE", "S_ALLOC_VGPR", "S_SLEEP_VAR", "S_GETPC_B64",
-        "S_SENDMSG_RTN_B32", "S_SENDMSG_RTN_B64"}
-# skip barriers, s_waits, wrap level atomics, and ray tracing (bvh)
-SKIP_SUBSTR = ["SAVEEXEC", "CMPX", "WREXEC", "MOVREL", "ATOMIC", "S_BUFFER_", "S_ATC_PROBE", "BARRIER", "S_WAITCNT", "BVH",
-               "DS_CMPSTORE_RTN", "DS_WRAP_RTN_B32", "DS_ORDERED_COUNT", "DS_GWS", "GS_REG", "GLOBAL_LOAD_LDS", "GLOBAL_STORE_BLOCK"]
-
-ALU_FORMATS = {"VOP1", "VOP1_LIT", "VOP1_SDST", "VOP2", "VOP2_LIT", "VOP3", "VOP3_SDST", "VOP3SD", "VOP3P", "VOP3P_MFMA", "VOP3PX2",
-               "VOPC", "SOP1", "SOP1_LIT", "SOP2", "SOP2_LIT", "SOPC", "SOPC_LIT", "SOPK", "SOPK_LIT", "VINTERP"}
-# intentionally not testing scratch memory ops
-MEM_FORMATS = {"VGLOBAL", "GLOBAL", "SMEM", "DS"}
-
-def should_skip(op:Enum) -> bool: return (name:=op.name) in SKIP or any(sub in name for sub in SKIP_SUBSTR)
-
-# ** named register assignments
-
-# ALU operands
-ALU_VGPR_STRIDE = 16          # v[0], v[16], v[32], ... per ALU operand slot
-ALU_SGPR_STRIDE = 4           # s[0], s[4], s[8], ... per ALU operand slot
-
-# memory address registers
-S_KERNARG_PTR = (0, 1)
-S_BUF_PTR = (2, 3)
-V_VADDR = (0, 1)
-V_DS_ADDR = 0
-
-# memory data registers
-MEM_VGPR_BASE = 32            # v[32], v[48], ... for vdst/vdata/vsrc
-MEM_VGPR_STRIDE = 16          # spacing between memory data vgpr slots
-MEM_SGPR_BASE = 8             # s[8], s[10], ... for SMEM sdata
-MEM_SGPR_STRIDE = 2           # spacing between memory data sgpr slots
-
-# ** create an ALU instruction based on the operands
-
-def create_alu_inst(op:Enum, builder:functools.partial[Inst]) -> Inst:
-  inst_cls, operands, slot = builder.func, OPERANDS[op], 0
-  kwargs:dict[str, Reg|int] = {}
-  for name, field in inst_cls._fields:
-    if isinstance(field, (FixedBitField, EnumBitField)): continue
-    nregs = max(1, operands[name][1] // 32) if name in operands else 1
-    is_sreg = name in operands and "SREG" in str(operands[name][2])
-    base_v, base_s = slot * ALU_VGPR_STRIDE, slot * ALU_SGPR_STRIDE
-    if name == "sdst" and isinstance(field, SGPRField): reg = VCC_LO
-    elif is_sreg and not isinstance(field, VGPRField): reg = VCC_LO
-    elif isinstance(field, VGPRField): reg = v[base_v:base_v+nregs-1] if nregs > 1 else v[base_v]
-    elif isinstance(field, SSrcField): reg = VCC_LO if nregs <= 2 else s[base_s:base_s+nregs-1] if nregs > 1 else s[base_s]
-    elif isinstance(field, SGPRField): reg = s[base_s:base_s+nregs-1] if nregs > 1 else s[base_s]
-    elif isinstance(field, SrcField): reg = v[base_v:base_v+nregs-1] if nregs > 1 else v[base_v]
-    else: reg = None
-    if reg is not None: kwargs[name] = reg; slot += 1
-    elif isinstance(field, BitField): kwargs[name] = field.default
-  return builder(**kwargs)
-
-# ** create a memory instruction with pre set address registers
-
-MEM_PRESET_REGS:dict[str, dict[str, Reg]] = {
-  "VGLOBAL":{"saddr":s[S_BUF_PTR[0]:S_BUF_PTR[1]], "vaddr":v[V_VADDR[0]:V_VADDR[1]]},
-  "GLOBAL":{"saddr":s[S_BUF_PTR[0]:S_BUF_PTR[1]], "addr":v[V_DS_ADDR]},  # addr is 32-bit offset when saddr is valid SGPR
-  "DS":{"addr":v[V_DS_ADDR]},
-  "SMEM":{"sbase":s[S_KERNARG_PTR[0]:S_KERNARG_PTR[1]], "soffset":NULL},
-}
-
-def create_mem_inst(op:Enum, builder:functools.partial[Inst]) -> Inst:
-  inst_cls, operands, field_map = builder.func, OPERANDS.get(op, {}), MEM_PRESET_REGS.get(builder.func.__name__, {})
-  kwargs:dict[str, Reg|int] = {}
-  vslot, sslot = 0, 0
-  for name, field in inst_cls._fields:
-    if isinstance(field, (FixedBitField, EnumBitField)): continue
-    if name in field_map:
-      kwargs[name] = field_map[name]
-      continue
-    nregs = max(1, operands[name][1] // 32) if name in operands else 1
-    if isinstance(field, VGPRField):
-      vi = MEM_VGPR_BASE + vslot * MEM_VGPR_STRIDE
-      kwargs[name] = v[vi:vi+nregs-1] if nregs > 1 else v[vi]
-      vslot += 1
-    elif isinstance(field, (SGPRField, AlignedSGPRField, SBaseField)):
-      si = MEM_SGPR_BASE + sslot * MEM_SGPR_STRIDE
-      kwargs[name] = s[si:si+nregs-1] if nregs > 1 else s[si]
-      sslot += 1
-    elif isinstance(field, BitField): kwargs[name] = field.default
-  return builder(**kwargs)
-
-# ** collect all memory and ALU instructions from the ISA autogen
-
-def collect_instructions() -> tuple[list[Inst], list[Inst], list[str]]:
-  op_map:dict[Enum, functools.partial[Inst]] = {}
-  for name, obj in inspect.getmembers(all_insts):
-    if isinstance(obj, functools.partial) and len(obj.args) == 1: op_map[obj.args[0]] = obj
-  alu_insts:list[Inst] = []
-  mem_insts:list[Inst] = []
-  skipped:list[str] = []
-  for op_enum, builder in op_map.items():
-    if should_skip(op_enum) or op_enum not in OPERANDS: skipped.append(op_enum.name); continue
-    fmt = builder.func.__name__
-    if fmt in ALU_FORMATS: alu_insts.append(create_alu_inst(op_enum, builder))
-    elif fmt in MEM_FORMATS: mem_insts.append(create_mem_inst(op_enum, builder))
-  return alu_insts, mem_insts, skipped
-
-def exec_insts(insts:list):
-  k = Kernel(arch)
-  # ** prologue for global memory
-  k.emit(s_load_b64(sdata=s[S_BUF_PTR[0]:S_BUF_PTR[1]], sbase=s[S_KERNARG_PTR[0]:S_KERNARG_PTR[1]], soffset=NULL))
-  k.waitcnt(lgkm=0)
-  k.emit(v_mov_b32_e32(v[V_VADDR[0]], 0))
-  k.emit(v_mov_b32_e32(v[V_VADDR[1]], 0))
-  # ** emit
-  for inst in insts: k.emit(inst)
-  k.emit(s_endpgm())
-  # ** run
-  NUM_THREADS, NUM_GRIDS, BUF_SIZE = 32, 1, 1024*1024
-  def fxn(A:UOp, B:UOp, C:UOp) -> UOp:
-    lidx, gidx = UOp.special(NUM_THREADS, "lidx0"), UOp.special(NUM_GRIDS, "gidx0")
-    lds = UOp(Ops.DEFINE_LOCAL, dtypes.uint8.ptr(size=BUF_SIZE, addrspace=AddrSpace.LOCAL), (), "lds")
-    sink = UOp.sink(A.base, B.base, C.base, lds, lidx, gidx, arg=KernelInfo(name="discover_ops"))
-    return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg="AMD"), UOp(Ops.LINEAR, src=tuple(UOp(Ops.INS, arg=x) for x in k.finalize()))))
-  A = Tensor.empty(BUF_SIZE, dtype=dtypes.uint8)
-  B = Tensor.empty(1, dtype=dtypes.uint8)
-  C = Tensor.empty(1, dtype=dtypes.uint8)
-  Tensor.custom_kernel(A, B, C, fxn=fxn)[0].realize()
-
-if __name__ == "__main__":
-  import sys
-  arch = Device[Device.DEFAULT].renderer.arch
-  if arch.startswith("gfx12"):
-    from tinygrad.runtime.autogen.amd.rdna4.ins import *
-    import tinygrad.runtime.autogen.amd.rdna4.ins as all_insts
-  elif arch.startswith("gfx11"):
-    from tinygrad.runtime.autogen.amd.rdna3.ins import *
-    import tinygrad.runtime.autogen.amd.rdna3.ins as all_insts
-    # these don"t exist in RDNA3, only RDNA3.5 and above
-    SKIP.update(["S_FMAAK_F32", "S_FMAMK_F32"])
-  else:
-    print(f"{arch} not supported yet")
-    sys.exit(0)
-  alu_insts, mem_insts, skipped = collect_instructions()
-  print(f"collected {len(alu_insts)} ALU + {len(mem_insts)} memory instructions ({len(skipped)} skipped)")
-  exec_insts(mem_insts+alu_insts)
@@ -9,7 +9,6 @@ EXAMPLES = [
  "test/backend/test_custom_kernel.py TestCustomKernel.test_empty",
  "test/test_tiny.py TestTiny.test_plus",
  "test/test_tiny.py TestTiny.test_gemm",
-  "extra/sqtt/examples/discover_ops.py"
 ]

 if __name__ == "__main__":
@@ -118,7 +118,7 @@ def decode(sqtt_evs:list[ProfileSQTTEvent], disasms:dict[str, dict[int, Inst]])
    nonlocal exc
    try: rocprof.rocprof_trace_decoder_parse_data(copy_cb, trace_cb, isa_cb, None)
    except AttributeError as e:
-      exc = RuntimeError("Failed to find rocprof-trace-decoder. Run sudo ./extra/sqtt/install_rocprof_decoder.py to install")
+      exc = RuntimeError("Failed to find rocprof-trace-decoder. Run sudo ./extra/sqtt/install_sqtt_decoder.py to install")
      exc.__cause__ = e
  (t:=threading.Thread(target=worker, daemon=True)).start()
  t.join()
@@ -1,6 +1,7 @@
 # simple tests
 import unittest
 import torch
+import warnings
 from tinygrad.helpers import getenv, GlobalCounters
 if getenv("TINY_BACKEND2"):
  import extra.torch_backend.backend2
@@ -17,7 +18,9 @@ class TestKernelFusionRegression(unittest.TestCase):
    torch.manual_seed(42)
    GlobalCounters.reset()
    fn().detach().cpu().numpy()
-    self.assertEqual(GlobalCounters.kernel_count, expected_kernels)
+    expectation = f"{GlobalCounters.kernel_count} vs {expected_kernels} expected."
+    if GlobalCounters.kernel_count < expected_kernels: warnings.warn(f"{expectation} Expectation can be lowered.", UserWarning)
+    self.assertLessEqual(GlobalCounters.kernel_count, expected_kernels, f"{expectation}")

  def test_elementwise_fusion(self):
    def fn():
@@ -31,7 +34,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, 6)
+    self._check_kernel_count(fn, 8)

  def test_batchnorm_fusion(self):
    def fn():
@@ -41,7 +44,7 @@ class TestKernelFusionRegression(unittest.TestCase):
      bn.eval()
      with torch.no_grad():
        return torch.nn.functional.relu(bn(conv(x)))
-    self._check_kernel_count(fn, 10)
+    self._check_kernel_count(fn, 16)

  def test_reduce_fusion(self):
    def fn():
@@ -89,7 +92,7 @@ class TestKernelFusionRegression(unittest.TestCase):
        out = bn(conv(x))
        out += identity
        return torch.nn.functional.relu(out)
-    self._check_kernel_count(fn, 12)
+    self._check_kernel_count(fn, 17)

  def test_multiple_inplace_ops_fusion(self):
    def fn():
@@ -114,7 +117,7 @@ class TestKernelFusionRegression(unittest.TestCase):
      bn.train()
      with torch.no_grad():
        return bn(x)
-    self._check_kernel_count(fn, 8)
+    self._check_kernel_count(fn, 10)

  # 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 +138,7 @@ class TestKernelFusionRegression(unittest.TestCase):
      loss.backward()
      optimizer.step()
      return loss
-    self._check_kernel_count(fn, 24)
+    self._check_kernel_count(fn, 28)

 if __name__ == "__main__":
  unittest.main()
@@ -1,17 +0,0 @@
-<?xml version="1.0" encoding="UTF-8"?>
-<!DOCTYPE plist PUBLIC "-//Apple//DTD PLIST 1.0//EN" "http://www.apple.com/DTDs/PropertyList-1.0.dtd">
-<plist version="1.0">
-<dict>
-	<key>com.apple.application-identifier</key>
-	<string>9YG3G8543N.org.tinygrad.tinygpu.edriver</string>
-	<key>com.apple.developer.driverkit</key>
-	<true/>
-	<key>com.apple.developer.driverkit.transport.pci</key>
-	<array>
-		<dict>
-			<key>IOPCIPrimaryMatch</key>
-			<string>0x000010de&amp;0x0000FFFF</string>
-		</dict>
-	</array>
-</dict>
-</plist>
@@ -1,33 +0,0 @@
-#!/bin/bash
-set -e
-
-xcodebuild clean build CODE_SIGN_IDENTITY="" CODE_SIGNING_REQUIRED=NO -alltargets -configuration Release build
-
-cp "../profiles/edriver_rel_2.provisionprofile" "./build/Release/TinyGPU.app/Contents/Library/SystemExtensions/org.tinygrad.tinygpu.edriver.dext/embedded.provisionprofile"
-cp "../profiles/installer_provisioning.provisionprofile" "./build/Release/TinyGPU.app/Contents/embedded.provisionprofile"
-
-codesign \
-    --sign "Developer ID Application: tinygrad, Corp. (9YG3G8543N)" \
-    --entitlements ./TinyGPUDriverExtension/TinyGPUDriver.NV.Release.entitlements \
-    --verbose \
-    --options runtime \
-    --timestamp \
-    --force \
-    ./build/Release/TinyGPU.app/Contents/Library/SystemExtensions/org.tinygrad.tinygpu.edriver.dext
-
-codesign \
-    --sign "Developer ID Application: tinygrad, Corp. (9YG3G8543N)" \
-    --entitlements ./macOS/macOS.entitlements \
-    --options runtime \
-    --verbose \
-    --timestamp \
-    --force \
-    ./build/Release/TinyGPU.app
-
-codesign --verify --deep --strict --verbose=4 ./build/Release/TinyGPU.app/Contents/Library/SystemExtensions/org.tinygrad.tinygpu.edriver.dext
-
-codesign --verify --deep --strict --verbose=4 ./build/Release/TinyGPU.app
-
-spctl -a -vv ./build/Release/TinyGPU.app
-
-spctl -a -vv ./build/Release/TinyGPU.app/Contents/Library/SystemExtensions/org.tinygrad.tinygpu.edriver.dext
@@ -1,17 +1,17 @@
 A command line tool for exploring the VIZ trace.

-After running with VIZ=-1, use `extra/viz/cli.py` to explore the saved trace files.
+After running with VIZ=-1, use `PYTHONPATH=. extra/viz/cli.py` to explore the saved trace files.

 ## Inspect runtime profiling

-Use `extra/viz/cli.py --profile` to list all traced devices.
+Use `PYTHONPATH=. extra/viz/cli.py --profile` to list all traced devices.

 List top slowest kernels on a device: `--profile --device "AMD"`
 List samples of a kernel on a device: `--profile --device "AMD" --kernel E_3`

 ## Inspect codegen and PatternMatcher

-Use `extra/viz/cli.py --rewrites` to list all traced kernels.
+Use `PYTHONPATH=. extra/viz/cli.py --rewrites` to list all traced kernels.

 List all codegen steps for a kernel: `--rewrites --kernel E_3`
 Get source code: `--rewrites --kernel E_3 --select "View Source"`
@@ -1,73 +1,44 @@
 #!/usr/bin/env python3
 import os
 os.environ["VIZ"] = "0"
-import argparse, pathlib, sys, struct, json
+import argparse, pathlib
 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
-
-# ** generic helpers
+from test.null.test_viz import load_profile

 def optional_eq(val:dict, arg:str|None) -> bool: return arg is None or ansistrip(val["name"]) == arg

 def print_data(data:dict) -> None:
  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(colored(line, "red" if line.startswith("-") else "green" if line.startswith("+") else None))
+      if m.get("uop"):
+        print("Input UOp:")
+        print(m["uop"])
+      if not m["diff"]: continue
+      print("Rewrites:")
+      fp = pathlib.Path(m["upat"][0][0])
+      print(f"{fp.parent.name}/{fp.name}:{m['upat'][0][1]}")
+      print(m["upat"][1])
+      for line in m["diff"]:
+        color = "red" if line.startswith("-") else "green" if line.startswith("+") else None
+        print(colored(line, color))
  if data.get("src") is not None: print(data["src"])

-# ** Profiler trace decoder
-
-# 0 means None, otherwise it's an enum value
-def option(i:int) -> int|None: return None if i == 0 else i-1
-
-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] = {}
-  for _ in range(layout_len):
-    klen = u("<B")[0]
-    k = ret[off:off+klen].decode()
-    off += klen
-    layout[k] = v = {"events":[]}
-    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["peak"] = u("<Q")[0]
-      for _ in range(event_count):
-        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}
-
 if __name__ == "__main__":
  parser = argparse.ArgumentParser()
  g_mode = parser.add_argument_group("mode")
  g_mode.add_argument("--profile", action="store_true", help="View profile trace")
  g_mode.add_argument("--rewrites", action="store_true", help="View rewrites trace")
-  g_common = parser.add_argument_group("common options")
-  g_common.add_argument("--kernel", type=str, default=None, metavar="NAME", help="Select a kernel by name (optional name, default: only list names)")
  g_profile = parser.add_argument_group("profile options")
  g_profile.add_argument("--device", type=str, default=None, metavar="NAME", help="Select a device (optional name, default: only list names)")
  g_profile.add_argument("--top", type=int, default=10, metavar="N", help="Number of top kernels to show (-1 for all, default: 10)")
  g_rewrites = parser.add_argument_group("rewrites options")
  g_rewrites.add_argument("--select", type=str, default=None, metavar="NAME",
                          help="Select an item within the chosen kernel (optional name, default: only list names)")
+  g_common = parser.add_argument_group("common options")
+  g_common.add_argument("--kernel", type=str, default=None, metavar="NAME", help="Select a kernel by name (optional name, default: only list names)")
  parser.add_argument("--profile-path", type=pathlib.Path, metavar="PATH", help="Path to profile (optional file, default: latest profile)",
                        default=pathlib.Path(temp("profile.pkl", append_user=True)))
  parser.add_argument("--rewrites-path", type=pathlib.Path, metavar="PATH", help="Path to rewrites (optional file, default: latest rewrites)",
@@ -75,14 +46,14 @@ if __name__ == "__main__":
  args = parser.parse_args()
  if not args.profile and not args.rewrites:
    parser.print_help()
-    sys.exit(0)
+    exit(0)

  viz.trace = viz.load_pickle(args.rewrites_path, default=RewriteTrace([], [], {}))
  viz.ctxs = viz.get_rewrites(viz.trace)

  if args.profile:
    from tabulate import tabulate
-    profile = decode_profile(viz.get_profile(viz.load_pickle(args.profile_path, default=[])))
+    profile = load_profile(viz.load_pickle(args.profile_path, default=[]))
    agg, total, n = {}, 0, 0
    if args.device is None: print("Select a device:")
    for k,v in profile["layout"].items():
@@ -92,7 +63,7 @@ if __name__ == "__main__":
      for e in v.get("events", []):
        et = e["dur"]*1e-6
        if args.kernel is not None:
-          if optional_eq(e, args.kernel) and n < 10:
+          if ansistrip(e["name"]) == args.kernel and n < 10:
            ptm = colored(time_to_str(et, w=9), "yellow" if et > 0.01 else None) if et is not None else ""
            name = e["name"]+(" " * (46 - ansilen(e["name"])))
            print(f"{name} {ptm}/{(et or 0)*1e3:9.2f}ms  "+e['fmt'].replace('\n', ' | ')+"  ")
@@ -110,7 +81,7 @@ if __name__ == "__main__":
        other_t = total-sum(t for _, (t, _) in sel)
        table.append([f"Other ({len(other)} unique)", time_to_str(other_t, w=9), sum(c for _,(_,c) in other), f"{other_t/total*100.0:.2f}%"])
      print(tabulate(table, headers=["name", "total", "count", "pct"], tablefmt="github"))
-    sys.exit(0)
+    exit(0)

  for k in viz.ctxs:
    if not optional_eq(k, args.kernel): continue
@@ -324,12 +324,6 @@ def _disasm_smem(inst: SMEM) -> str:
  if name in ('s_memrealtime', 's_memtime'): return f"{name} {_fmt_sdst(inst.sdata, dst_n, cdna)}"
  return f"{name} {_fmt_sdst(inst.sdata, dst_n, cdna)}, {sbase_str}, {off_s}" + _mods((inst.glc, " glc"), (getattr(inst, 'dlc', 0), " dlc"))

-R4_TH_LOAD = {1: 'TH_LOAD_NT', 2: 'TH_LOAD_HT', 3: 'TH_LOAD_LU', 4: 'TH_LOAD_RT_WB', 5: 'TH_LOAD_NT_WB'}
-R4_TH_STORE = {1: 'TH_STORE_NT', 2: 'TH_STORE_HT', 3: 'TH_STORE_ST', 4: 'TH_STORE_RT_WB', 5: 'TH_STORE_NT_WB'}
-R4_TH_ATOMIC = {1: 'TH_ATOMIC_RETURN', 2: 'TH_ATOMIC_NT', 3: 'TH_ATOMIC_RETURN_NT',
-                4: 'TH_ATOMIC_CASCADE_RT', 5: 'TH_ATOMIC_CASCADE_RETURN', 6: 'TH_ATOMIC_CASCADE_NT', 7: 'TH_ATOMIC_CASCADE_RETURN_NT'}
-R4_SCOPE = {1: 'SCOPE_SE', 2: 'SCOPE_DEV', 3: 'SCOPE_SYS'}
-
 def _disasm_flat(inst: FLAT) -> str:
  name, cdna, r4 = inst.op_name.lower(), _is_cdna(inst), _is_r4(inst)
  acc = getattr(inst, 'acc', 0)
@@ -337,10 +331,9 @@ def _disasm_flat(inst: FLAT) -> str:
  if r4: seg = 'flat' if (cls_name:=inst.__class__.__name__) == 'VFLAT' else ('global' if cls_name == 'VGLOBAL' else 'scratch')
  else: seg = ['flat', 'scratch', 'global'][inst.seg] if inst.seg < 3 else 'flat'
  instr = f"{seg}_{name.split('_', 1)[1] if '_' in name else name}"
-  # Global/scratch uses 13-bit signed offset (RDNA3/CDNA), 24-bit signed offset (RDNA4)
+  # Global/scratch uses 13-bit signed offset
  offset = inst.ioffset if r4 else inst.offset  # type: ignore[attr-defined]
-  if r4: off_val = offset if offset < (1 << 23) else offset - (1 << 24)  # sign extend 24-bit
-  elif seg != 'flat':
+  if seg != 'flat':
    if cdna:
      # CDNA: bit 12 is sign bit but not in offset field
      raw = int.from_bytes(inst.to_bytes(), 'little')
@@ -355,9 +348,7 @@ def _disasm_flat(inst: FLAT) -> str:
  w = regs.get('data', regs.get('d', 1)) if 'store' in name or 'atomic' in name else regs.get('d', 1)
  off_s = f" offset:{off_val}" if off_val else ""
  if cdna: mods = f"{off_s}{' sc0' if inst.sc0 else ''}{' nt' if inst.nt else ''}{' sc1' if getattr(inst, 'sc1', 0) else ''}"  # type: ignore[attr-defined]
-  elif r4:
-    th_names = R4_TH_ATOMIC if 'atomic' in name else (R4_TH_STORE if 'store' in name else R4_TH_LOAD)
-    mods = off_s + (f" th:{th_names[inst.th]}" if inst.th in th_names else "") + (f" scope:{R4_SCOPE[inst.scope]}" if inst.scope in R4_SCOPE else "")
+  elif r4: mods = f"{off_s}{' scope' if inst.scope else ''}{' th' if inst.th else ''}"  # type: ignore[attr-defined]
  else: mods = f"{off_s}{' glc' if inst.glc else ''}{' slc' if inst.slc else ''}{' dlc' if inst.dlc else ''}"
  if seg == 'flat': saddr_s = ""
  elif _unwrap(inst.saddr) in (0x7F, 124): saddr_s = ", off"
@@ -366,7 +357,7 @@ def _disasm_flat(inst: FLAT) -> str:
    saddr_s = f", {(SPECIAL_PAIRS_CDNA if cdna else SPECIAL_PAIRS)[_unwrap(inst.saddr)]}"
  elif t := _ttmp(inst.saddr, 2): saddr_s = f", {t}"
  else: saddr_s = f", {_sreg(inst.saddr, 2) if _unwrap(inst.saddr) < 106 else decode_src(_unwrap(inst.saddr), cdna)}"
-  if 'addtid' in name: return f"{instr} {reg_fn((inst.vsrc if r4 else inst.data) if 'store' in name else inst.vdst)}{saddr_s}{mods}"
+  if 'addtid' in name: return f"{instr} {reg_fn(inst.data if 'store' in name else inst.vdst)}{saddr_s}{mods}"
   # RDNA4: vaddr instead of addr, vsrc instead of data
  addr = inst.vaddr if r4 else inst.addr  # type: ignore[attr-defined]
  data = inst.vsrc if r4 else inst.data  # type: ignore[attr-defined]
@@ -381,7 +372,7 @@ def _disasm_flat(inst: FLAT) -> str:
  addr_s = "off" if not inst.sve and seg == 'scratch' else _vreg(addr, addr_w)
  data_s, vdst_s = reg_fn(data, w), reg_fn(inst.vdst, w // 2 if 'cmpswap' in name else w)
  if 'atomic' in name:
-    glc_or_sc0 = inst.sc0 if cdna else (inst.th & 1 if r4 else inst.glc)  # type: ignore[attr-defined]
+    glc_or_sc0 = inst.sc0 if cdna else inst.glc  # type: ignore[attr-defined]
    sfx = f"{saddr_s if seg != 'flat' else ''}{mods}"
    return f"{instr} {vdst_s}, {addr_s}, {data_s}{sfx}" if glc_or_sc0 else f"{instr} {addr_s}, {data_s}{sfx}"
  if 'store' in name: return f"{instr} {addr_s}, {data_s}{saddr_s}{mods}"
@@ -104,34 +104,6 @@ class TestCmpClass(unittest.TestCase):
    st = run_program(instructions, n_lanes=1)
    self.assertEqual(st.vcc & 1, 0, "Signaling NaN should not match quiet mask")

-  def test_v_cmp_lg_f32_nan(self):
-    """v_cmp_lg_f32 is ordered not-equal (<>): NaN <> x should be False per IEEE 754."""
-    quiet_nan = 0x7fc00000
-    one_f32 = 0x3f800000  # 1.0f
-    instructions = [
-      s_mov_b32(s[0], quiet_nan),
-      v_mov_b32_e32(v[0], s[0]),
-      s_mov_b32(s[1], one_f32),
-      v_mov_b32_e32(v[1], s[1]),
-      v_cmp_lg_f32_e32(v[0], v[1]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vcc & 1, 0, "v_cmp_lg_f32(NaN, 1.0) should be 0")
-
-  def test_v_cmp_neq_f32_nan(self):
-    """v_cmp_neq_f32 is unordered not-equal (!=): NaN != x should be True per IEEE 754."""
-    quiet_nan = 0x7fc00000
-    one_f32 = 0x3f800000  # 1.0f
-    instructions = [
-      s_mov_b32(s[0], quiet_nan),
-      v_mov_b32_e32(v[0], s[0]),
-      s_mov_b32(s[1], one_f32),
-      v_mov_b32_e32(v[1], s[1]),
-      v_cmp_neq_f32_e32(v[0], v[1]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vcc & 1, 1, "v_cmp_neq_f32(NaN, 1.0) should be 1")
-
  def test_v_cmp_sets_vcc_bits(self):
    """V_CMP_EQ sets VCC bits based on per-lane comparison."""
    instructions = [
@@ -40,7 +40,7 @@ RDNA4_FILES = ['gfx12_asm_sop1.s', 'gfx12_asm_sop2.s', 'gfx12_asm_sopp.s', 'gfx1
  'gfx12_asm_vop1.s', 'gfx12_asm_vop2.s', 'gfx12_asm_vopc.s', 'gfx12_asm_vopcx.s', 'gfx12_asm_vop3.s', 'gfx12_asm_vop3c.s',
  'gfx12_asm_vop3cx.s', 'gfx12_asm_vop3p.s', 'gfx12_asm_vop3_from_vop1.s', 'gfx12_asm_vop3_from_vop2.s',
  'gfx12_asm_vop3p_features.s', 'gfx12_asm_vopd.s', 'gfx12_asm_vopd_features.s',
-  'gfx12_asm_ds.s', 'gfx12_asm_smem.s', 'gfx12_asm_vflat.s',
+  'gfx12_asm_ds.s', 'gfx12_asm_smem.s',
  'gfx12_asm_wmma_w32.s']

 def _parse_llvm_tests(text: str, pattern: str) -> list[tuple[str, bytes]]:
@@ -21,7 +21,7 @@ OTHER_SIMD_OPS = {InstOp.OTHER_LDS_LOAD, InstOp.OTHER_LDS_STORE, InstOp.OTHER_LD
                  InstOp.OTHER_FLAT_STORE_128, InstOp.OTHER_GLOBAL_LOAD, InstOp.OTHER_GLOBAL_LOAD_VADDR,
                  InstOp.OTHER_GLOBAL_STORE_64, InstOp.OTHER_GLOBAL_STORE_96, InstOp.OTHER_GLOBAL_STORE_128,
                  InstOp.OTHER_GLOBAL_STORE_VADDR_128}
-OTHER_SIMD_OPS_RDNA4 = {InstOpRDNA4.OTHER_VMEM}
+OTHER_SIMD_OPS_RDNA4 = {InstOpRDNA4.OTHER_VMEM, InstOpRDNA4.OTHER_VMEM_STORE}

 # ═══════════════════════════════════════════════════════════════════════════════
 # ROCPROF DECODER
@@ -0,0 +1,183 @@
+"""Tests comparing sqtt.py PACKET_TYPES_RDNA3/RDNA4 against AMD's rocprof-trace-decoder binary."""
+import unittest, struct, ctypes, pickle
+from pathlib import Path
+
+ROCPROF_LIB = Path("/usr/lib/librocprof-trace-decoder.so")
+import tinygrad
+EXAMPLES_DIR = Path(tinygrad.__file__).parent.parent / "extra/sqtt/examples"
+
+# CDNA pkt_fmt -> size in bytes (extracted from rocprof hash table)
+CDNA_PKT_SIZES = {0: 2, 1: 8, 2: 8, 3: 4, 4: 2, 5: 6, 6: 2, 7: 2, 8: 2, 9: 2, 10: 2, 11: 8, 12: 6, 13: 4, 14: 8, 15: 6}
+
+def _find_segment(perms: str):
+  """Find a segment of the loaded library with given permissions (e.g. 'rw-p', 'r--p')."""
+  with open('/proc/self/maps', 'r') as f:
+    for line in f:
+      if 'librocprof-trace-decoder.so' in line and f' {perms} ' in line:
+        parts = line.split()
+        return int(parts[0].split('-')[0], 16), int(parts[2], 16)
+  return None, None
+
+def _read_array(file_offset: int, count: int):
+  """Read an array of uint8 at file_offset from the loaded library."""
+  base, seg_offset = _find_segment('rw-p')
+  if base is None: return None
+  return list((ctypes.c_uint8 * count).from_address(base + (file_offset - seg_offset)))
+
+def _load_lib():
+  if not ROCPROF_LIB.exists(): return False
+  ctypes.CDLL(str(ROCPROF_LIB))
+  return True
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# RDNA EXTRACTION (nibble-based format)
+# ═══════════════════════════════════════════════════════════════════════════════
+
+def extract_bit_tables():
+  """Extract bit budget tables. Returns (layout2, layout3, layout4) or None."""
+  if not _load_lib(): return None
+  return _read_array(0x2d220, 32), _read_array(0x2d280, 32), _read_array(0x2d2c0, 32)
+
+def extract_delta_fields():
+  """Extract delta bitfield tables. Returns (layout2, layout3, layout4) dicts mapping type_id -> (lo, hi)."""
+  if not _load_lib(): return None
+  ro_base, ro_offset = _find_segment('r--p')
+  if ro_base is None: return None
+
+  def read_table(file_offset, num_entries):
+    addr = ro_base + (file_offset - ro_offset)
+    data = bytes((ctypes.c_uint8 * (num_entries * 12)).from_address(addr))
+    return {type_id: (lo, hi) for j in range(0, len(data), 12)
+            for type_id, lo, hi in [struct.unpack('<III', data[j:j+12])] if type_id < 32}
+
+  return read_table(0x26800, 24), read_table(0x26dc0, 25), read_table(0x27300, 27)
+
+def extract_packet_encodings():
+  """Extract packet encodings. Returns (L2, L3, L4) dicts mapping type_id -> (mask, value)."""
+  if not _load_lib(): return None
+  rw_base, rw_offset = _find_segment('rw-p')
+  if rw_base is None: return None
+
+  # Read base encodings from registration vector at 0x2d340
+  vec_start = ctypes.c_void_p.from_address(rw_base + (0x2d340 - rw_offset)).value
+  vec_end = ctypes.c_void_p.from_address(rw_base + (0x2d348 - rw_offset)).value
+  base = {}
+  if vec_start and vec_end:
+    for i in range((vec_end - vec_start) // 32):
+      addr = vec_start + i * 32
+      type_id = ctypes.c_uint8.from_address(addr).value
+      pat_start = ctypes.c_void_p.from_address(addr + 8).value
+      pat_end = ctypes.c_void_p.from_address(addr + 16).value
+      if pat_start and pat_end and 0 < (n := pat_end - pat_start) <= 8:
+        pat = list((ctypes.c_uint8 * n).from_address(pat_start))
+        base[type_id] = (sum(1 << j for j in range(n)), sum(b << j for j, b in enumerate(pat)))
+
+  return {**base, 17: (0x7f, 0x51), 25: (0x7f, 0x31)}, base, {**base}  # L2 has overrides
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# CDNA EXTRACTION (16-bit header format)
+# ═══════════════════════════════════════════════════════════════════════════════
+
+def extract_cdna_packet_sizes():
+  """Extract CDNA pkt_fmt -> size mapping by running rocprof decoder to populate its hash table."""
+  if not _load_lib(): return None
+  from test.amd.test_sqtt_examples import run_rocprof_decoder
+
+  if not (pkl_path := next((EXAMPLES_DIR / "gfx950").glob("*.pkl"), None)): return None
+  with open(pkl_path, "rb") as f: data = pickle.load(f)
+  sqtt_events = [e for e in data if type(e).__name__ == "ProfileSQTTEvent"]
+  prg = next((e for e in data if type(e).__name__ == "ProfileProgramEvent"), None)
+  if not sqtt_events or not prg: return None
+
+  # Run decoder to trigger hash table initialization
+  run_rocprof_decoder([e.blob for e in sqtt_events], prg.lib, prg.base, "gfx950")
+
+  # Extract hash table: head at 0x2d4f0, nodes are 16 bytes (next[8], key[4], value[4])
+  rw_base, rw_offset = _find_segment('rw-p')
+  if not (head := ctypes.c_void_p.from_address(rw_base + (0x2d4f0 - rw_offset)).value if rw_base else None): return None
+
+  pkt_sizes: dict[int, int] = {}
+  node, seen = head, set()
+  while node and node not in seen and len(pkt_sizes) < 20:
+    seen.add(node)
+    key, val = ctypes.c_uint32.from_address(node + 8).value, ctypes.c_uint32.from_address(node + 12).value
+    if key < 16 and val in (0x10, 0x20, 0x30, 0x40): pkt_sizes[key] = {0x10: 2, 0x20: 4, 0x30: 6, 0x40: 8}[val]
+    node = ctypes.c_void_p.from_address(node).value  # type: ignore[assignment]
+  return pkt_sizes if len(pkt_sizes) == 16 else None
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# TESTS
+# ═══════════════════════════════════════════════════════════════════════════════
+
+class TestSQTTMatchesBinary(unittest.TestCase):
+  def test_bit_counts_match_layout3(self): self._test_bit_counts(3)
+  def test_bit_counts_match_layout4(self): self._test_bit_counts(4)
+  def test_encodings_match_layout3(self): self._test_encodings(3)
+  def test_encodings_match_layout4(self): self._test_encodings(4)
+  def test_delta_fields_match_layout3(self): self._test_delta_fields(3)
+  def test_delta_fields_match_layout4(self): self._test_delta_fields(4)
+
+  def test_cdna_packet_sizes(self):
+    """Extract and verify CDNA pkt_fmt -> size mapping from rocprof's hash table."""
+    if not (EXAMPLES_DIR / "gfx950").exists(): self.skipTest("no CDNA examples")
+    if not (pkt_sizes := extract_cdna_packet_sizes()): self.skipTest("rocprof-trace-decoder not installed")
+    for pkt_fmt, size in CDNA_PKT_SIZES.items():
+      with self.subTest(pkt_fmt=pkt_fmt): self.assertEqual(pkt_sizes.get(pkt_fmt), size)
+
+  def test_cdna_packet_definitions(self):
+    from tinygrad.renderer.amd.sqtt import PACKET_TYPES_CDNA
+    for pkt_fmt, pkt_cls in PACKET_TYPES_CDNA.items():
+      with self.subTest(packet=pkt_cls.__name__):
+        self.assertEqual(pkt_cls.encoding.default, pkt_fmt)
+        self.assertEqual(CDNA_PKT_SIZES[pkt_fmt] * 2, pkt_cls._size_nibbles)  # type: ignore[attr-defined]
+
+  def _test_bit_counts(self, layout: int):
+    if not (tables := extract_bit_tables()): self.skipTest("rocprof-trace-decoder not installed")
+    from tinygrad.renderer.amd.sqtt import PACKET_TYPES_RDNA3, PACKET_TYPES_RDNA4
+    for type_id, pkt_cls in {3: PACKET_TYPES_RDNA3, 4: PACKET_TYPES_RDNA4}[layout].items():
+      with self.subTest(packet=pkt_cls.__name__):
+        self.assertEqual(pkt_cls._size_nibbles * 4, tables[layout - 2][type_id])  # type: ignore[attr-defined]
+
+  def _test_encodings(self, layout: int):
+    if not (encodings := extract_packet_encodings()): self.skipTest("rocprof-trace-decoder not installed")
+    from tinygrad.renderer.amd.sqtt import PACKET_TYPES_RDNA3, PACKET_TYPES_RDNA4
+    for type_id, pkt_cls in {3: PACKET_TYPES_RDNA3, 4: PACKET_TYPES_RDNA4}[layout].items():
+      with self.subTest(packet=pkt_cls.__name__):
+        self.assertEqual((pkt_cls.encoding.mask, pkt_cls.encoding.default), encodings[layout - 2][type_id])
+
+  def _test_delta_fields(self, layout: int):
+    if not (deltas := extract_delta_fields()): self.skipTest("rocprof-trace-decoder not installed")
+    from tinygrad.renderer.amd.sqtt import PACKET_TYPES_RDNA3, PACKET_TYPES_RDNA4
+    for type_id, pkt_cls in {3: PACKET_TYPES_RDNA3, 4: PACKET_TYPES_RDNA4}[layout].items():
+      if type_id not in deltas[layout - 2]: continue
+      delta = getattr(pkt_cls, 'delta', None)
+      actual = (0, 0) if delta is None else (delta.lo, delta.hi + 1)
+      with self.subTest(packet=pkt_cls.__name__): self.assertEqual(actual, deltas[layout - 2][type_id])
+
+if __name__ == "__main__":
+  tables = extract_bit_tables()
+  encodings = extract_packet_encodings()
+  deltas = extract_delta_fields()
+
+  TYPE_NAMES = {1: 'VALUINST', 2: 'VMEMEXEC', 3: 'ALUEXEC', 4: 'IMMEDIATE', 5: 'IMMEDIATE_MASK', 6: 'WAVERDY',
+    7: 'TS_DELTA_S8_W3', 8: 'WAVEEND', 9: 'WAVESTART', 10: 'TS_DELTA_S5_W2', 11: 'WAVEALLOC', 12: 'TS_DELTA_S5_W3',
+    13: 'PERF', 14: 'UTILCTR', 15: 'TS_DELTA_SHORT', 16: 'NOP', 17: 'TS_WAVE_STATE', 18: 'EVENT', 19: 'EVENT_BIG',
+    20: 'REG', 21: 'SNAPSHOT', 22: 'TS_DELTA_OR_MARK', 23: 'LAYOUT_HEADER', 24: 'INST', 25: 'UNK_25'}
+
+  print("L2:", tables[0], "\nL3:", tables[1], "\nL4:", tables[2])
+  if encodings and tables:
+    print(f"\n{'TypeID':>6} {'Name':>18} {'L2 enc':>12} {'L3 enc':>12} {'L4 enc':>12}"
+          f" {'L2':>4} {'L3':>4} {'L4':>4} {'L2 delta':>12} {'L3 delta':>12} {'L4 delta':>12}")
+    print("-" * 140)
+    for type_id in sorted(set(encodings[0]) | set(encodings[1]) | set(encodings[2])):
+      name = TYPE_NAMES.get(type_id, f'UNK_{type_id}')
+      bits = [tables[i][type_id] if type_id < len(tables[i]) else 0 for i in range(3)]
+      enc_strs = [f"0x{encodings[i][type_id][0]:02x}/0x{encodings[i][type_id][1]:02x}" if type_id in encodings[i] else "-" for i in range(3)]
+      delta_strs = [f"[{d[1]-1}:{d[0]}]" if (d := deltas[i].get(type_id, (0, 0)))[1] > d[0] else "-" for i in range(3)]
+      print(f"{type_id:6d} {name:>18} {enc_strs[0]:>12} {enc_strs[1]:>12} {enc_strs[2]:>12}"
+            f" {bits[0]:4d} {bits[1]:4d} {bits[2]:4d} {delta_strs[0]:>12} {delta_strs[1]:>12} {delta_strs[2]:>12}")
+
+  cdna = extract_cdna_packet_sizes()
+  if cdna: print(f"\nCDNA packet sizes: {cdna}")
+
+  unittest.main()
@@ -2,7 +2,7 @@
 import unittest, pickle
 from typing import Iterator
 from pathlib import Path
-from tinygrad.helpers import DEBUG, OSX
+from tinygrad.helpers import DEBUG
 from tinygrad.renderer.amd.sqtt import print_packets, map_insts
 from tinygrad.runtime.autogen.amd.rdna3.ins import s_endpgm
 from test.amd.disasm import disasm
@@ -10,7 +10,7 @@ from test.amd.disasm import disasm
 import tinygrad
 EXAMPLES_DIR = Path(tinygrad.__file__).parent.parent / "extra/sqtt/examples"

-def rocprof_inst_traces_match(sqtt, prg, target, pass_rocprof_err=False):
+def rocprof_inst_traces_match(sqtt, prg, target):
  from tinygrad.viz.serve import amd_decode
  from extra.sqtt.roc import decode as roc_decode, InstExec
  addr_table = amd_decode(prg.lib, target)
@@ -24,13 +24,13 @@ def rocprof_inst_traces_match(sqtt, prg, target, pass_rocprof_err=False):

  passed_insts = 0
  for pkt, info in map_insts(sqtt.blob, prg.lib, target):
-    if DEBUG >= 2: print_packets([(pkt, info)])
+    if DEBUG >= 2: print_packets([pkt])
    if info is None: continue
    if DEBUG >= 2: print(f"{' '*29}{disasm(info.inst)}")
    rocprof_inst = next(rwaves_iter[info.wave][0])
    ref_pc = rocprof_inst.pc-prg.base
    # always check pc matches
-    assert ref_pc == info.pc or pass_rocprof_err, f"pc mismatch {ref_pc}:{disasm_map[rocprof_inst.pc]} != {info.pc}:{disasm(info.inst)}"
+    assert ref_pc == info.pc, f"pc mismatch {ref_pc}:{disasm_map[rocprof_inst.pc]} != {info.pc}:{disasm(info.inst)}"
    # special handling for s_endpgm, it marks the wave completion.
    if info.inst == s_endpgm():
      completed_wave = list(rwaves_iter[info.wave].pop(0))
@@ -67,9 +67,7 @@ class TestSQTTMapBase(unittest.TestCase):
        if not event.itrace: continue
        if event.kern not in kern_events: continue
        with self.subTest(example=name, kern=event.kern):
-          # rocprof OSX has a bug for sopk decoding, linux rocprof works
-          pass_rocprof_err = OSX and target == "gfx1200" and name.startswith("profile_py")
-          passed_insts, n_waves, n_units = rocprof_inst_traces_match(event, kern_events[event.kern], target, pass_rocprof_err)
+          passed_insts, n_waves, n_units = rocprof_inst_traces_match(event, kern_events[event.kern], target)
          if n_waves: print(f"{name}: passed for {passed_insts} instructions across {n_waves} waves scheduled on {n_units} wave units")

 class TestSQTTMapRDNA3(TestSQTTMapBase): target = "gfx1100"
@@ -47,18 +47,6 @@ def verify_asm_gemm(batch:int, M:int, N:int, K:int, dtype=dtypes.float16, gpus:i
 def verify_asm_gemm_k_sharded(M:int, N:int, K:int, dtype=dtypes.float16, gpus:int=8) -> None:
  run_asm_gemm((M, K), (K, N), dtype=dtype, a_shard=1, b_shard=0, gpus=gpus)

-def verify_asm_gemm_n_sharded(batch:int, M:int, N:int, K:int, dtype=dtypes.float16, gpus:int=2) -> None:
-  run_asm_gemm((batch, M, K), (K, N), dtype=dtype, a_shard=None, b_shard=1, gpus=gpus)
-
-def verify_asm_gemm_m_sharded(M:int, N:int, K:int, dtype=dtypes.float16, gpus:int=2) -> None:
-  run_asm_gemm((M, K), (K, N), dtype=dtype, a_shard=0, b_shard=None, gpus=gpus)
-
-def verify_asm_gemm_n_sharded_2d(M:int, N:int, K:int, dtype=dtypes.float16, gpus:int=2) -> None:
-  run_asm_gemm((M, K), (K, N), dtype=dtype, a_shard=None, b_shard=1, gpus=gpus)
-
-def verify_asm_gemm_k_sharded_3d(batch:int, M:int, N:int, K:int, dtype=dtypes.float16, gpus:int=2) -> None:
-  run_asm_gemm((batch, M, K), (K, N), dtype=dtype, a_shard=2, b_shard=0, gpus=gpus)
-
 # 128x smaller than usual
 # uses the UOp GEMM, runs on non CDNA4 and CI
@unittest.skipUnless(is_dtype_supported(dtypes.half), "need half")
@@ -72,14 +60,6 @@ class TestGemm(unittest.TestCase):
  def test_gemm_multi(self): verify_asm_gemm(2, 64, 32, 32, gpus=2)
  @needs_second_gpu
  def test_gemm_k_sharded(self): verify_asm_gemm_k_sharded(64, 64, 2*64, gpus=2)
-  @needs_second_gpu
-  def test_gemm_m_sharded(self): verify_asm_gemm_m_sharded(2*64, 64, 32, gpus=2)
-  @needs_second_gpu
-  def test_gemm_n_sharded(self): verify_asm_gemm_n_sharded(1, 64, 64, 32, gpus=2)
-  @needs_second_gpu
-  def test_gemm_n_sharded_2d(self): verify_asm_gemm_n_sharded_2d(64, 2*64, 32, gpus=2)
-  @needs_second_gpu
-  def test_gemm_k_sharded_3d(self): verify_asm_gemm_k_sharded_3d(1, 64, 32, 2*64, gpus=2)

 # uses the Asm GEMM on CDNA4 only for speed reasons
 class TestGemmLarge(unittest.TestCase):
@@ -121,20 +101,6 @@ class TestGemmLarge(unittest.TestCase):
      verify_asm_gemm(3, 256, 256, 256)
  def test_gemm_previously_unsupported(self): verify_asm_gemm(8, 1024, 1024, 4096, gpus=8)

-  # M-sharded 2D
-  def test_m_sharded_1(self): verify_asm_gemm_m_sharded(8*8192, 4096, 4096, dtype=dtypes.bfloat16, gpus=8)
-  def test_m_sharded_2(self): verify_asm_gemm_m_sharded(8*4096, 14336, 4096, dtype=dtypes.bfloat16, gpus=8)
-
-  # N-sharded 2D
-  def test_n_sharded_2d_1(self): verify_asm_gemm_n_sharded_2d(8192, 8*4096, 4096, dtype=dtypes.bfloat16, gpus=8)
-  def test_n_sharded_2d_2(self): verify_asm_gemm_n_sharded_2d(4096, 8*14336, 4096, dtype=dtypes.bfloat16, gpus=8)
-
-  # tensor parallel shapes (Llama 8B, MP=8)
-  def test_tp_n_sharded_wq(self): verify_asm_gemm_n_sharded(1, 8192, 4096, 4096, dtype=dtypes.bfloat16, gpus=8)
-  def test_tp_n_sharded_w1(self): verify_asm_gemm_n_sharded(1, 8192, 14336, 4096, dtype=dtypes.bfloat16, gpus=8)
-  def test_tp_k_sharded_wo(self): verify_asm_gemm_k_sharded_3d(1, 8192, 4096, 4096, dtype=dtypes.bfloat16, gpus=8)
-  def test_tp_k_sharded_w2(self): verify_asm_gemm_k_sharded_3d(1, 8192, 4096, 14336, dtype=dtypes.bfloat16, gpus=8)
-
  # more shapes: vary M, N, K independently
  def test_shape_small_square(self): verify_asm_gemm(1, 256, 256, 256)
  def test_shape_small_rect_m(self): verify_asm_gemm(1, 512, 256, 256)
@@ -27,11 +27,6 @@ class TestMovedConstFolding(unittest.TestCase):
  def test_add_padded_one(self):
    _check_ast_count(1, Tensor([1.0, 2, 3, 4]) * Tensor.ones(2).pad(((1, 1),)))

-  def test_copy_padded_const(self):
-    schedule = Tensor.ones(4, device="CPU:0").pad(((1, 1),)).to("CPU:1").schedule()
-    assert not any(si.ast.op is Ops.COPY for si in schedule), "const copy should be folded"
-    np.testing.assert_equal(Tensor.ones(4, device="CPU:0").pad(((1, 1),)).to("CPU:1").numpy(), [0, 1, 1, 1, 1, 0])
-
  def test_cast_padded(self):
    # NOTE: it's always 1 kernel when calling .numpy, limitation of _check_ast_count
    if is_dtype_supported(dtypes.int16):
@@ -228,17 +228,17 @@ class TestMultiTensor(unittest.TestCase):
      a,b = _test_allreduce(Tensor.rand(256, 256))
      np.testing.assert_almost_equal(a.numpy(), b.numpy(), decimal=5)

-  def test_multiple_to_single_device(self):
-    kernel_counts = {}
-    for ring in (0, 2):
-      GlobalCounters.reset()
-      with Context(RING=ring, SCACHE=0):
-        t = Tensor.arange(32).contiguous().shard(devices_4, 0).to(Device.DEFAULT)
-        t.realize()
-      kernel_counts[ring] = GlobalCounters.kernel_count
-      self.assertEqual(t.device, Device.DEFAULT)
-      np.testing.assert_equal(t.numpy(), np.arange(32))
-    self.assertNotEqual(kernel_counts[0], kernel_counts[2])
+  def test_multiple_to_single_device_naive(self):
+    with Context(RING=0):
+      t = Tensor.arange(32).shard(devices_4, 0).to(Device.DEFAULT).realize()
+    self.assertEqual(t.device, Device.DEFAULT)
+    np.testing.assert_equal(t.numpy(), np.arange(32))
+
+  def test_multiple_to_single_device_ring(self):
+    with Context(RING=2):
+      t = Tensor.arange(32).shard(devices_4, 0).to(Device.DEFAULT).realize()
+    self.assertEqual(t.device, Device.DEFAULT)
+    np.testing.assert_equal(t.numpy(), np.arange(32))

  def test_allreduce_all2all(self):
    with Context(ALL2ALL=2):
@@ -795,38 +795,6 @@ class TestSchedule(unittest.TestCase):
      self.assertIsNotNone(out.uop.base.realized)
      self.assertIsInstance(out.uop.base.realized.dtype, ImageDType)

-  @unittest.skipIf(Device.DEFAULT != "CL", "image only supported on CL")
-  def test_image_dot_f16_fusion(self):
-    with Context(FLOAT16=1):
-      def cnt():
-        x, y, z = Tensor.empty((64, 64), dtype='float'), Tensor.empty((64, 64), dtype='float'), Tensor.empty((64, 64), dtype='float')
-        a = (x @ y).relu()
-        sched = ((a @ z).relu() + a).schedule()
-        for si in sched: si.lower()
-        return len([si for si in sched if isinstance(si.prg, CompiledRunner)])
-
-      with Context(IMAGE=1): cnt1 = cnt()
-      with Context(IMAGE=2): cnt2 = cnt()
-
-      self.assertEqual(cnt1, 5)
-      self.assertEqual(cnt2, 5)
-
-  @unittest.skipIf(Device.DEFAULT != "CL", "image only supported on CL")
-  @unittest.expectedFailure
-  def test_image_conv_fusion(self):
-    def cnt():
-      x, y, z = Tensor.empty((1, 4, 3, 3)), Tensor.empty((4, 1, 3, 3)), Tensor.empty((4, 1, 7, 7))
-      a = x.conv2d(y, Tensor.empty(4), groups=4, padding=1)
-      b = a.conv2d(z, groups=4, padding=3)
-      sched = (a + b).schedule()
-      for si in sched: si.lower()
-      return len([si for si in sched if isinstance(si.prg, CompiledRunner)])
-
-    with Context(IMAGE=1): cnt1 = cnt()
-    with Context(IMAGE=2): cnt2 = cnt()
-
-    self.assertEqual(cnt1, cnt2)
-
  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))
@@ -205,20 +205,6 @@ class TestSetitem(unittest.TestCase):
    n[:, ind_1.numpy(), :, ind_2.numpy(), :] = v.numpy()
    np.testing.assert_equal(t.numpy(), n)

-  def test_setitem_tensor_int_indexing(self):
-    t = Tensor.zeros(4, 3, dtype=dtypes.int).contiguous()
-    t[Tensor([0, 2]), 0] = Tensor([99, 88], dtype=dtypes.int)
-    n = np.zeros((4, 3), dtype=np.int32)
-    n[[0, 2], 0] = [99, 88]
-    np.testing.assert_equal(t.numpy(), n)
-
-  def test_setitem_tensor_slice_indexing(self):
-    t = Tensor.zeros(4, 3, dtype=dtypes.int).contiguous()
-    t[Tensor([0, 2]), :2] = Tensor([[10, 20], [30, 40]], dtype=dtypes.int)
-    n = np.zeros((4, 3), dtype=np.int32)
-    n[[0, 2], :2] = [[10, 20], [30, 40]]
-    np.testing.assert_equal(t.numpy(), n)
-
  def test_setitem_2d_tensor_indexing(self):
    t = Tensor.zeros(2, dtype=dtypes.int).contiguous()
    index = Tensor([[0, 1], [1,0]])
@@ -293,43 +279,17 @@ class TestWithGrad(unittest.TestCase):
    x = Tensor.rand(8)
    z[:3] = x

+  def test_set_into_requires_grad(self):
+    z = Tensor.rand(8, 8, requires_grad=True)
+    x = Tensor.rand(8)
+    with self.assertRaises(NotImplementedError):
+      z[:3] = x
+
  def test_set_with_requires_grad(self):
-    z = Tensor.ones(8, 8)
-    x = Tensor.rand(8, 8, requires_grad=True)
-    z[:] = x
-    z.sum().backward()
-    np.testing.assert_allclose(x.grad.numpy(), np.ones((8, 8)))
-
-  def test_set_nonleaf_requires_grad(self):
-    x = Tensor([1.0, 2.0, 3.0, 4.0], requires_grad=True)
-    z = x * 2
-    z[:2] = Tensor([10.0, 20.0])
-    z.sum().backward()
-    np.testing.assert_allclose(x.grad.numpy(), [0, 0, 2, 2])
-
-  def test_set_overlapping_requires_grad(self):
-    z = Tensor.zeros(6, requires_grad=True)
-    x = Tensor.ones(4, requires_grad=True)
-    y = Tensor.ones(4, requires_grad=True) * 2
-    z[:4] = x
-    z[2:] = y
-    z.sum().backward()
-    np.testing.assert_allclose(x.grad.numpy(), [1, 1, 0, 0])
-    np.testing.assert_allclose(y.grad.numpy(), np.ones(4))
-
-  def test_set_iadd_requires_grad(self):
-    z = Tensor([1.0, 2.0, 3.0, 4.0], requires_grad=True)
-    x = Tensor([10.0, 20.0], requires_grad=True)
-    z[:2] += x
-    z.sum().backward()
-    np.testing.assert_allclose(z.grad.numpy(), np.ones(4))
-    np.testing.assert_allclose(x.grad.numpy(), np.ones(2))
-
-  def test_set_used_before_setitem(self):
-    z = Tensor([1.0, 2.0, 3.0, 4.0], requires_grad=True)
-    _ = z.sum()
-    with self.assertRaises(RuntimeError):
-      z[:2] = Tensor([0.0, 0.0])
+    z = Tensor.rand(8, 8)
+    x = Tensor.rand(8, requires_grad=True)
+    with self.assertRaises(NotImplementedError):
+      z[:3] = x

 class TestSetitemLoop(unittest.TestCase):
  def test_arange(self):
@@ -25,7 +25,7 @@ class TestStunning(unittest.TestCase):
    nv = a[12].cat(a[76]).tolist()

    vi = Variable('i', 0, a.shape[0]-1)
-    with self.assertRaisesRegex(RuntimeError, "bind mismatch on"):
+    with self.assertRaisesRegex(AssertionError, "bind mismatch on"):
      wv = a[vi.bind(12)].cat(a[vi.bind(76)]).tolist()
      self.assertListEqual(nv, wv)

@@ -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
 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"}, "only NV, AMD, CUDA")
  def test_subbuffer_transfer(self):
    t = Tensor.arange(0, 10, dtype=dtypes.uint8).realize()
    vt = t[2:5].contiguous().realize()
@@ -69,58 +69,6 @@ class TestSymbolicOps(unittest.TestCase):
        # symbolic shape dropout is not supported
        self.test_attention(dropout_p=0.5)

-  def test_sdpa_symbolic_seq_len(self):
-    # symbolic seq_len on all of q/k/v (dim -2 after transpose)
-    q = Tensor.rand(2, 10, 4, 8)
-    k = Tensor.rand(2, 10, 4, 8)
-    v = Tensor.rand(2, 10, 4, 8)
-    for i in range(1, 5):
-      vi = Variable("i", 1, 10).bind(i)
-      Tensor.realize(q, k, v)
-      symbolic = q[:, :vi].transpose(1, 2).scaled_dot_product_attention(
-        k[:, :vi].transpose(1, 2), v[:, :vi].transpose(1, 2)).realize()[:2, :4, :i, :8].numpy()
-      expected = q[:, :i].transpose(1, 2).scaled_dot_product_attention(
-        k[:, :i].transpose(1, 2), v[:, :i].transpose(1, 2)).realize().numpy()
-      np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
-
-  def test_sdpa_symbolic_seq_len_query_only(self):
-    # symbolic seq_len on query only (dim -2 after transpose)
-    q = Tensor.rand(2, 10, 4, 8)
-    k = Tensor.rand(2, 5, 4, 8)
-    v = Tensor.rand(2, 5, 4, 8)
-    for i in range(1, 5):
-      vi = Variable("i", 1, 10).bind(i)
-      Tensor.realize(q, k, v)
-      symbolic = q[:, :vi].transpose(1, 2).scaled_dot_product_attention(
-        k.transpose(1, 2), v.transpose(1, 2)).realize()[:2, :4, :i, :8].numpy()
-      expected = q[:, :i].transpose(1, 2).scaled_dot_product_attention(
-        k.transpose(1, 2), v.transpose(1, 2)).realize().numpy()
-      np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
-
-  def test_sdpa_symbolic_batch(self):
-    # symbolic batch dim (dim 0)
-    q = Tensor.rand(10, 4, 3, 8)
-    k = Tensor.rand(10, 4, 3, 8)
-    v = Tensor.rand(10, 4, 3, 8)
-    for i in range(1, 5):
-      vi = Variable("i", 1, 10).bind(i)
-      Tensor.realize(q, k, v)
-      symbolic = q[:vi].scaled_dot_product_attention(k[:vi], v[:vi]).realize()[:i, :4, :3, :8].numpy()
-      expected = q[:i].scaled_dot_product_attention(k[:i], v[:i]).realize().numpy()
-      np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
-
-  def test_sdpa_symbolic_heads(self):
-    # symbolic heads dim (dim -3)
-    q = Tensor.rand(2, 10, 3, 8)
-    k = Tensor.rand(2, 10, 3, 8)
-    v = Tensor.rand(2, 10, 3, 8)
-    for i in range(1, 5):
-      vi = Variable("i", 1, 10).bind(i)
-      Tensor.realize(q, k, v)
-      symbolic = q[:, :vi].scaled_dot_product_attention(k[:, :vi], v[:, :vi]).realize()[:2, :i, :3, :8].numpy()
-      expected = q[:, :i].scaled_dot_product_attention(k[:, :i], v[:, :i]).realize().numpy()
-      np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
-
  def test_attention_pos_0_sz_0(self):
    Attention(128, 8)(Tensor.ones(1, 0, 128), Variable("start_pos", 0, 128).bind(0), None)

@@ -136,30 +136,6 @@ class TestTensorVariable(unittest.TestCase):
    with self.assertRaises(AssertionError):
      t.chunk(2, dim=0)  # can't split along symbolic dim

-  def test_symbolic_var_sum(self, var_name="u"):
-    t = Variable("t", 1, 10).bind(4)
-    v = Variable(var_name, 1, 5).bind(1)
-    mask = (Tensor.full((1, 1, t, v+t), 1) + 1).contiguous()
-    mask.shrink(((0, 1), (0, 1), (0, 4), (0, 4))).numpy()
-  def test_symbolic_var_sum_alt_name(self): self.test_symbolic_var_sum("s")
-
-  def test_symbolic_triu(self):
-    t = Variable("t", 1, 10).bind(4)
-    for start_pos in (0, 1, 3):
-      var_start_pos = Variable("start_pos", 0, 5).bind(start_pos)
-      mask = Tensor.full((1, 1, t, var_start_pos+t), float("-inf")).triu(var_start_pos+1)
-      out = mask.shrink(((0, 1), (0, 1), (0, 4), (0, start_pos+4))).numpy()
-      expected = np.triu(np.full((1, 1, 4, start_pos+4), float("-inf")), k=start_pos+1)
-      np.testing.assert_equal(out, expected)
-
-  def test_symbolic_tril(self):
-    t = Variable("t", 1, 10).bind(4)
-    for start_pos in (0, 1, 3):
-      var_start_pos = Variable("start_pos", 0, 5).bind(start_pos)
-      mask = Tensor.full((1, 1, t, var_start_pos+t), float("-inf")).tril(var_start_pos+1)
-      out = mask.shrink(((0, 1), (0, 1), (0, 4), (0, start_pos+4))).numpy()
-      expected = np.tril(np.full((1, 1, 4, start_pos+4), float("-inf")), k=start_pos+1)
-      np.testing.assert_equal(out, expected)

 if __name__ == '__main__':
  unittest.main()
@@ -1,31 +0,0 @@
-#!/usr/bin/env python3
-"""
-Stress test for beam timeout + device recovery on AM devices.
-
-Usage:
-  AMD=1 python test/external/external_test_beam_timeout_recovery.py
-"""
-from tinygrad import Tensor, Device
-from tinygrad.helpers import Context
-from tinygrad.runtime.ops_amd import AMDDevice
-
-if __name__ == "__main__":
-  dev = Device["AMD"]
-  assert isinstance(dev, AMDDevice) and dev.is_am(), "not am"
-
-  N = 10000
-  for i in range(N):
-    with Context(DEBUG=0, BEAM=0):
-      a = Tensor.rand(4096, 4096, device="AMD").contiguous().realize()
-      b = Tensor.rand(4096, 4096, device="AMD").contiguous().realize()
-      c = a.matmul(b)
-      c.realize()
-    try: dev.synchronize(timeout=1)
-    except RuntimeError as e: print(e)
-    with Context(DEBUG=0, BEAM=0):
-      a = Tensor.ones(512, 512, device="AMD").contiguous().realize()
-      b = Tensor.ones(512, 512, device="AMD").contiguous().realize()
-      result = a.matmul(b).realize()[0, 0].item()
-    assert result == 512.0, f"iter {i}: got {result}"
-    print(f"  iter {i+1}/{N}: ok")
-  print(f"=== All {N} iterations passed ===")
@@ -1,55 +0,0 @@
-import subprocess, sys, os, random
-
-CHILD_SCRIPT = """
-import os, random
-import numpy as np
-from tinygrad import Tensor, Device
-from tinygrad.runtime.ops_amd import AMDDevice
-
-dev = Device["AMD"]
-for i in range({N}):
-  sz = random.randint(1, {MAX_SZ})
-  data = np.random.randint(0, 256, sz, dtype=np.uint8)
-  t = Tensor(data, device="AMD").contiguous().realize()
-  dev.synchronize()
-  result = t.numpy()
-  assert (result == data).all(), f"Data mismatch at iter {{i}}"
-""".strip()
-
-def run_child(n_ops, max_sz, timeout):
-  env = os.environ.copy()
-  env.setdefault("SDMA_RING_SIZE", "4096")
-
-  script = CHILD_SCRIPT.format(N=n_ops, MAX_SZ=max_sz)
-  p = subprocess.Popen([sys.executable, "-c", script], stdout=subprocess.PIPE, stderr=subprocess.PIPE, env=env)
-
-  try:
-    _, stderr = p.communicate(timeout=timeout)
-    return ("ok" if p.returncode == 0 else "fail"), stderr.decode(errors='replace')
-  except subprocess.TimeoutExpired:
-    p.kill()
-    p.communicate()
-    return "timeout", "TIMEOUT: SDMA ring likely stuck"
-
-if __name__ == "__main__":
-  n_iters = int(os.environ.get("FUZZ_ITERS", "10000"))
-  timeout = int(os.environ.get("FUZZ_TIMEOUT", "10"))
-  max_sz = int(os.environ.get("FUZZ_MAX_SZ", "65536"))
-
-  timeouts = 0
-  failures = 0
-
-  for i in range(n_iters):
-    # Run child with many ops to stress the small sdma ring buffer across warm starts
-    n_ops = random.randint(20, 100)
-    status, stderr = run_child(n_ops=n_ops, max_sz=max_sz, timeout=timeout)
-    if status == "timeout":
-      timeouts += 1
-      print(f"\tstderr: {stderr[:500]}")
-    elif status == "fail":
-      failures += 1
-      print(f"\tstderr: {stderr[:500]}")
-    else:
-      print(f"iter {i}: ok (n_ops={n_ops})")
-
-  print(f"\n=== Results: {n_iters} iterations, {timeouts} timeouts, {failures} failures ===")
@@ -4,19 +4,13 @@
 These tests intentionally cause GPU faults to verify error handling.
 Run with: AMD=1 python -m pytest test/external/external_test_gpu_crash.py -v
 """
-import unittest, re, importlib
+import unittest, re
 from tinygrad.device import Device
+from tinygrad.runtime.autogen.amd.rdna3.ins import *  # noqa: F403
 from tinygrad.renderer.amd.dsl import s, v, Inst, NULL

-RDNA3_CDNA3_MAP = {"v_mov_b32_e32": "v_mov_b32_e32", "s_mov_b32": "s_mov_b32", "s_waitcnt": "s_waitcnt", "s_endpgm": "s_endpgm",
-                   "global_load_b32": "global_load_dword", "global_store_b32": "global_store_dword",
-                   "global_atomic_add_u32": "global_atomic_add", "flat_load_b32": "flat_load_dword",
-                   "flat_store_b32": "flat_store_dword", "flat_atomic_add_u32": "flat_atomic_add", "s_load_b32": "s_load_dword"}
-
-def assemble(code:str, name:str="test", is_cdna:bool=False) -> str:
-  kd = {"next_free_vgpr": 8, "next_free_sgpr": 8, "user_sgpr_kernarg_segment_ptr": 1, "kernarg_size": 8}
-  if is_cdna: kd["accum_offset"] = 8
-  else: kd["wavefront_size32"] = 1
+def assemble(code:str, name:str="test") -> str:
+  kd = {"next_free_vgpr": 8, "next_free_sgpr": 8, "wavefront_size32": 1, "user_sgpr_kernarg_segment_ptr": 1, "kernarg_size": 8}
  return f".text\n.globl {name}\n.p2align 8\n.type {name},@function\n{name}:\n{code}\n.rodata\n.p2align 6\n.amdhsa_kernel {name}\n" + \
         "\n".join(f".amdhsa_{k} {v}" for k,v in kd.items()) + "\n.end_amdhsa_kernel"

@@ -27,10 +21,6 @@ class TestGPUCrash(unittest.TestCase):
    from tinygrad.runtime.support.compiler_amd import HIPCompiler
    cls.dev = Device["AMD"]
    cls.compiler = HIPCompiler(cls.dev.arch)
-    cls.is_cdna = cls.dev.target[0] < 10
-    ins = importlib.import_module('tinygrad.runtime.autogen.amd.' + ('cdna' if cls.is_cdna else 'rdna3') + '.ins')
-    for rdna3_name, cdna3_name in RDNA3_CDNA3_MAP.items():
-      setattr(cls, rdna3_name, getattr(ins, cdna3_name if cls.is_cdna else rdna3_name))

  def setUp(self):
    # Verify device works before each test
@@ -43,7 +33,7 @@ class TestGPUCrash(unittest.TestCase):

  def _run(self, code: str):
    from tinygrad.runtime.ops_amd import AMDProgram
-    prg = AMDProgram(self.dev, "test", self.compiler.compile(assemble(code, is_cdna=self.is_cdna)))
+    prg = AMDProgram(self.dev, "test", self.compiler.compile(assemble(code)))
    prg(self.dev.allocator.alloc(64), global_size=(1,1,1), local_size=(1,1,1), wait=True)

  def _run_insts(self, insts: list[Inst]):
@@ -67,32 +57,32 @@ class TestOutOfBoundsMemoryAccess(TestGPUCrash):

  def test_global_load_null_ptr(self):
    """Global load from NULL pointer."""
-    insts = [self.v_mov_b32_e32(v[0], 0), self.v_mov_b32_e32(v[1], 0),
-             self.global_load_b32(v[2], addr=v[0:1], saddr=NULL, offset=0), self.s_waitcnt(0), self.s_endpgm()]
+    insts = [v_mov_b32_e32(v[0], 0), v_mov_b32_e32(v[1], 0),
+             global_load_b32(v[2], addr=v[0:1], saddr=NULL, offset=0), s_waitcnt(0), s_endpgm()]
    self._assert_gpu_fault(lambda: self._run_insts(insts))

  def test_global_store_null_ptr(self):
    """Global store to NULL pointer."""
-    insts = [self.v_mov_b32_e32(v[0], 0), self.v_mov_b32_e32(v[1], 0), self.v_mov_b32_e32(v[2], 0xDEADBEEF),
-             self.global_store_b32(addr=v[0:1], data=v[2], saddr=NULL, offset=0), self.s_waitcnt(0), self.s_endpgm()]
+    insts = [v_mov_b32_e32(v[0], 0), v_mov_b32_e32(v[1], 0), v_mov_b32_e32(v[2], 0xDEADBEEF),
+             global_store_b32(addr=v[0:1], data=v[2], saddr=NULL, offset=0), s_waitcnt(0), s_endpgm()]
    self._assert_gpu_fault(lambda: self._run_insts(insts))

  def test_global_load_unmapped_high_address(self):
    """Global load from high unmapped address (0xDEAD00000000)."""
-    insts = [self.v_mov_b32_e32(v[0], 0x00000000), self.v_mov_b32_e32(v[1], 0xDEAD),
-             self.global_load_b32(v[2], addr=v[0:1], saddr=NULL, offset=0), self.s_waitcnt(0), self.s_endpgm()]
+    insts = [v_mov_b32_e32(v[0], 0x00000000), v_mov_b32_e32(v[1], 0xDEAD),
+             global_load_b32(v[2], addr=v[0:1], saddr=NULL, offset=0), s_waitcnt(0), s_endpgm()]
    self._assert_gpu_fault(lambda: self._run_insts(insts))

  def test_global_store_unmapped_high_address(self):
    """Global store to high unmapped address."""
-    insts = [self.v_mov_b32_e32(v[0], 0x00000000), self.v_mov_b32_e32(v[1], 0xDEAD), self.v_mov_b32_e32(v[2], 0x12345678),
-             self.global_store_b32(addr=v[0:1], data=v[2], saddr=NULL, offset=0), self.s_waitcnt(0), self.s_endpgm()]
+    insts = [v_mov_b32_e32(v[0], 0x00000000), v_mov_b32_e32(v[1], 0xDEAD), v_mov_b32_e32(v[2], 0x12345678),
+             global_store_b32(addr=v[0:1], data=v[2], saddr=NULL, offset=0), s_waitcnt(0), s_endpgm()]
    self._assert_gpu_fault(lambda: self._run_insts(insts))

  def test_global_atomic_unmapped(self):
    """Atomic operation on unmapped memory."""
-    insts = [self.v_mov_b32_e32(v[0], 0xBEEF0000), self.v_mov_b32_e32(v[1], 0xDEAD), self.v_mov_b32_e32(v[2], 1),
-             self.global_atomic_add_u32(addr=v[0:1], data=v[2], saddr=NULL, offset=0), self.s_waitcnt(0), self.s_endpgm()]
+    insts = [v_mov_b32_e32(v[0], 0xBEEF0000), v_mov_b32_e32(v[1], 0xDEAD), v_mov_b32_e32(v[2], 1),
+             global_atomic_add_u32(addr=v[0:1], data=v[2], saddr=NULL, offset=0), s_waitcnt(0), s_endpgm()]
    self._assert_gpu_fault(lambda: self._run_insts(insts))


@@ -101,14 +91,14 @@ class TestSMEMFaults(TestGPUCrash):

  def test_smem_load_null(self):
    """SMEM load from NULL base."""
-    insts = [self.s_mov_b32(s[2], 0), self.s_mov_b32(s[3], 0),
-             self.s_load_b32(s[4], s[2:3], 0, soffset=NULL), self.s_waitcnt(0), self.s_endpgm()]
+    insts = [s_mov_b32(s[2], 0), s_mov_b32(s[3], 0),
+             s_load_b32(s[4], s[2:3], 0, soffset=NULL), s_waitcnt(0), s_endpgm()]
    self._assert_gpu_fault(lambda: self._run_insts(insts))

  def test_smem_load_unmapped(self):
    """SMEM load from unmapped address."""
-    insts = [self.s_mov_b32(s[2], 0xBEEF0000), self.s_mov_b32(s[3], 0xDEAD),
-             self.s_load_b32(s[4], s[2:3], 0, soffset=NULL), self.s_waitcnt(0), self.s_endpgm()]
+    insts = [s_mov_b32(s[2], 0xBEEF0000), s_mov_b32(s[3], 0xDEAD),
+             s_load_b32(s[4], s[2:3], 0, soffset=NULL), s_waitcnt(0), s_endpgm()]
    self._assert_gpu_fault(lambda: self._run_insts(insts))


@@ -117,20 +107,20 @@ class TestFlatMemoryFaults(TestGPUCrash):

  def test_flat_load_null(self):
    """FLAT load from NULL address."""
-    insts = [self.v_mov_b32_e32(v[0], 0), self.v_mov_b32_e32(v[1], 0),
-             self.flat_load_b32(v[2], addr=v[0:1], saddr=NULL, offset=0), self.s_waitcnt(0), self.s_endpgm()]
+    insts = [v_mov_b32_e32(v[0], 0), v_mov_b32_e32(v[1], 0),
+             flat_load_b32(v[2], addr=v[0:1], saddr=NULL, offset=0), s_waitcnt(0), s_endpgm()]
    self._assert_gpu_fault(lambda: self._run_insts(insts))

  def test_flat_store_null(self):
    """FLAT store to NULL address."""
-    insts = [self.v_mov_b32_e32(v[0], 0), self.v_mov_b32_e32(v[1], 0), self.v_mov_b32_e32(v[2], 0xDEADBEEF),
-             self.flat_store_b32(addr=v[0:1], data=v[2], saddr=NULL, offset=0), self.s_waitcnt(0), self.s_endpgm()]
+    insts = [v_mov_b32_e32(v[0], 0), v_mov_b32_e32(v[1], 0), v_mov_b32_e32(v[2], 0xDEADBEEF),
+             flat_store_b32(addr=v[0:1], data=v[2], saddr=NULL, offset=0), s_waitcnt(0), s_endpgm()]
    self._assert_gpu_fault(lambda: self._run_insts(insts))

  def test_flat_atomic_null(self):
    """FLAT atomic on NULL address."""
-    insts = [self.v_mov_b32_e32(v[0], 0), self.v_mov_b32_e32(v[1], 0), self.v_mov_b32_e32(v[2], 1),
-             self.flat_atomic_add_u32(addr=v[0:1], data=v[2], saddr=NULL, offset=0), self.s_waitcnt(0), self.s_endpgm()]
+    insts = [v_mov_b32_e32(v[0], 0), v_mov_b32_e32(v[1], 0), v_mov_b32_e32(v[2], 1),
+             flat_atomic_add_u32(addr=v[0:1], data=v[2], saddr=NULL, offset=0), s_waitcnt(0), s_endpgm()]
    self._assert_gpu_fault(lambda: self._run_insts(insts))


@@ -2,9 +2,8 @@
 import subprocess, sys
 from tinygrad.helpers import getenv

-LOOPS = getenv("LOOPS", 50)
+LOOPS = getenv("LOOPS", 10)
 BROKEN = getenv("BROKEN", 0)
-ONLY_RESET = getenv("ONLY_RESET", 0)

 BROKEN_KERNEL_SCRIPT = """
 from tinygrad.device import Device
@@ -37,7 +36,7 @@ for i in range(LOOPS):
    print(f"=== Running broken kernel ({i+1}/{LOOPS}) ===")
    ret = subprocess.run([sys.executable, "-c", BROKEN_KERNEL_SCRIPT])
    print(f"=== broken kernel exited with code {ret.returncode} ===")
-  elif not ONLY_RESET:
-    print(f"=== Running test_tiny.py ({i+1}/{LOOPS}) ===")
-    ret = subprocess.run([sys.executable, "test/test_tiny.py", "TestTiny.test_plus"])
-    print(f"=== test_tiny.py exited with code {ret.returncode} ===")
+
+  print(f"=== Running test_tiny.py ({i+1}/{LOOPS}) ===")
+  ret = subprocess.run([sys.executable, "test/test_tiny.py", "TestTiny.test_plus"])
+  print(f"=== test_tiny.py exited with code {ret.returncode} ===")
@@ -10,7 +10,7 @@ from tinygrad.helpers import Profiling

 class FakeProgram:
  def __init__(self, name:str, prg:bytes, **kwargs): pass
-  def __call__(self, *bufs, global_size, local_size, vals=(), wait=False, **kw): pass
+  def __call__(self, *bufs, global_size, local_size, vals=(), wait=False): pass

 class FakeAllocator(Allocator[Compiled]):
  def _alloc(self, sz, options): return None
@@ -416,10 +416,10 @@ class Parser:
      case '||' | '|': return left | right
      case '&&' | '&': return left & right
      case '^': return left ^ right
-      case '==': return left.eq(right)
+      case '==' | '<>': return left.eq(right) if op == '==' else left.ne(right)
      case '!=': return left.ne(right)
-      case '>=' | '<=' | '>' | '<' | '<>':
-        ops = {'>=':(lambda a,b:a>=b),'<=':(lambda a,b:a<=b),'>':(lambda a,b:a>b),'<':(lambda a,b:a<b),'<>':(lambda a,b:a.ne(b))}
+      case '>=' | '<=' | '>' | '<':
+        ops = {'>=':(lambda a,b:a>=b),'<=':(lambda a,b:a<=b),'>':(lambda a,b:a>b),'<':(lambda a,b:a<b)}
        return self._cmp_nan(left, right, ops[op])
      case '>>' | '<<': return (left >> right) if op == '>>' else (left << right)
      case '+' | '-':
@@ -87,7 +87,7 @@ class TestHuggingFaceOnnxModels(unittest.TestCase):
      "input_ids": np.random.randint(0, 250002, (1, 11), dtype=np.int64),
      "attention_mask": np.ones((1, 11), dtype=np.int64),
    }
-    self._validate(repo_id, model_file, custom_inputs, atol=1e-3)
+    self._validate(repo_id, model_file, custom_inputs)

 if __name__ == "__main__":
  unittest.main()
@@ -349,184 +349,5 @@ class TestStopEarly(unittest.TestCase):
    ret = (c+d).substitute({c:cn}, extra_pm=pm_cvisit)
    assert ret == cn+d

-class TestWalkRewrite(unittest.TestCase):
-  """Tests for graph_rewrite with walk=True (MLIR Walk Pattern Rewrite Driver semantics).
-  walk=True gives a single-pass traversal that does NOT revisit or re-traverse into rewritten subtrees.
-  Supports both top-down (default) and bottom-up (bottom_up=True) modes."""
-
-  # *** top-down walk (default): process children first, then try pm on rebuilt node ***
-
-  def test_walk_topdown_simple_substitute(self):
-    a = UOp.variable('a', 0, 10)
-    b = UOp.variable('b', 0, 10)
-    ret = graph_rewrite(a + 4, _substitute, {a:b}, walk=True)
-    self.assertIs(ret, b+4)
-
-  def test_walk_topdown_does_not_traverse_into_replacement(self):
-    """Top-down walk: replacement subtrees are NOT re-entered."""
-    a = UOp.variable('a', 0, 10)
-    b = UOp.variable('b', 0, 10)
-    c = UOp.variable('c', 0, 10)
-    d = UOp.variable('d', 0, 10)
-    # a is replaced by b+c, but b inside the replacement is NOT further substituted to d
-    ret_walk = graph_rewrite(a + 4, _substitute, {a:b+c, b:d}, walk=True)
-    self.assertIs(ret_walk, (b+c)+4)
-    # contrast: greedy bottom_up WOULD replace b inside the replacement
-    ret_greedy = graph_rewrite(a + 4, _substitute, {a:b+c, b:d}, bottom_up=True)
-    self.assertIs(ret_greedy, (d+c)+4)
-
-  def test_walk_topdown_no_fixed_point(self):
-    """A bouncing pattern applies once and stops instead of looping."""
-    a = UOp.const(dtypes.int, 3)
-    pm = PatternMatcher([
-      (UPat(Ops.CONST, arg=3, name="x"), lambda x: x.replace(arg=4)),
-      (UPat(Ops.CONST, arg=4, name="x"), lambda x: x.replace(arg=3)),
-    ])
-    with self.assertRaises(RuntimeError):
-      graph_rewrite(a, pm, bottom_up=True)
-    ret = graph_rewrite(a, pm, walk=True)
-    self.assertIs(ret, UOp.const(dtypes.int, 4))
-
-  def test_walk_topdown_rewrites_children(self):
-    a = UOp.variable('a', 0, 10)
-    b = UOp.variable('b', 0, 10)
-    c = UOp.variable('c', 0, 10)
-    ret = graph_rewrite((a + 4) + (b + 5), _substitute, {a:c, b:c}, walk=True)
-    self.assertIs(ret, (c + 4) + (c + 5))
-
-  def test_walk_topdown_diamond(self):
-    a = UOp.variable('a', 0, 10)
-    b = UOp.variable('b', 0, 10)
-    ret = graph_rewrite((a + 4) + (a + 5), _substitute, {a:b}, walk=True)
-    self.assertIs(ret, (b + 4) + (b + 5))
-
-  def test_walk_topdown_children_rewritten_before_parent(self):
-    """Top-down walk processes children first: child substitution changes the rebuilt parent."""
-    a = UOp.variable('a', 0, 10, dtype=dtypes.float)
-    n1 = a.sin()          # sin(a)
-    ret = n1.sin()         # sin(sin(a))
-    # sin(a)->sqrt(a) fires first (child), parent rebuilds to sin(sqrt(a)), which doesn't match sin(sin(a)) in dvars
-    ret_walk = graph_rewrite(ret, _substitute, {a.sin():a.sqrt(), n1.sin():n1.sqrt()}, walk=True)
-    self.assertIs(ret_walk, a.sqrt().sin())
-
-  def test_walk_topdown_self_referential_replacement(self):
-    """Replacement containing the replaced node works without infinite recursion."""
-    a = UOp.variable('a', 0, 10, dtype=dtypes.float)
-    ret = graph_rewrite(a.sin() + 4, _substitute, {a.sin(): a.sin().sqrt()}, walk=True)
-    self.assertIs(ret, a.sin().sqrt() + 4)
-
-  def test_walk_topdown_visit_order(self):
-    """Top-down walk fires pm after children are processed (post-order)."""
-    visited = []
-    def track_visit(ctx, x):
-      ctx.append(x.arg if x.op is Ops.CONST else x.op)
-      return None
-    pm = PatternMatcher([(UPat(GroupOp.All, name="x"), track_visit)])
-    a = UOp.const(dtypes.int, 1)
-    b = UOp.const(dtypes.int, 2)
-    graph_rewrite(a + b, pm, ctx=visited, walk=True)
-    self.assertEqual(visited, [1, 2, Ops.ADD])
-
-  # *** bottom-up walk: try bpm on node first, skip children if it matches ***
-
-  def test_walk_bottomup_simple_substitute(self):
-    a = UOp.variable('a', 0, 10)
-    b = UOp.variable('b', 0, 10)
-    ret = graph_rewrite(a + 4, _substitute, {a:b}, bottom_up=True, walk=True)
-    self.assertIs(ret, b+4)
-
-  def test_walk_bottomup_does_not_traverse_into_replacement(self):
-    """Bottom-up walk: replacement subtrees are NOT entered."""
-    a = UOp.variable('a', 0, 10)
-    b = UOp.variable('b', 0, 10)
-    c = UOp.variable('c', 0, 10)
-    d = UOp.variable('d', 0, 10)
-    ret = graph_rewrite(a + 4, _substitute, {a:b+c, b:d}, bottom_up=True, walk=True)
-    self.assertIs(ret, (b+c)+4)
-
-  def test_walk_bottomup_parent_match_skips_children(self):
-    """Bottom-up walk matches parent first: if it matches, children are never visited."""
-    a = UOp.variable('a', 0, 10, dtype=dtypes.float)
-    n1 = a.sin()
-    ret = n1.sin()         # sin(sin(a))
-    # sin(sin(a)) matches n1.sin()->n1.sqrt() immediately, children never visited, sin(a) inside replacement untouched
-    ret_walk = graph_rewrite(ret, _substitute, {a.sin():a.sqrt(), n1.sin():n1.sqrt()}, bottom_up=True, walk=True)
-    self.assertIs(ret_walk, a.sin().sqrt())
-
-  def test_walk_bottomup_no_fixed_point(self):
-    """Bottom-up walk also applies once per node, no fixed-point iteration."""
-    a = UOp.const(dtypes.int, 3)
-    pm = PatternMatcher([
-      (UPat(Ops.CONST, arg=3, name="x"), lambda x: x.replace(arg=4)),
-      (UPat(Ops.CONST, arg=4, name="x"), lambda x: x.replace(arg=3)),
-    ])
-    ret = graph_rewrite(a, pm, bottom_up=True, walk=True)
-    self.assertIs(ret, UOp.const(dtypes.int, 4))
-
-  def test_walk_bottomup_visit_order(self):
-    """Bottom-up walk fires bpm before descending (pre-order)."""
-    visited = []
-    def track_visit(ctx, x):
-      ctx.append(x.arg if x.op is Ops.CONST else x.op)
-      return None
-    pm = PatternMatcher([(UPat(GroupOp.All, name="x"), track_visit)])
-    a = UOp.const(dtypes.int, 1)
-    b = UOp.const(dtypes.int, 2)
-    graph_rewrite(a + b, pm, ctx=visited, bottom_up=True, walk=True)
-    # bpm fires on each node before children: +, 1, 2
-    self.assertEqual(visited, [Ops.ADD, 1, 2])
-
-  def test_walk_bottomup_unmatched_falls_through_to_children(self):
-    """Bottom-up walk: if bpm doesn't match a node, its children are still processed."""
-    a = UOp.variable('a', 0, 10)
-    b = UOp.variable('b', 0, 10)
-    c = UOp.variable('c', 0, 10)
-    # only a is in dvars, not a+4. bpm won't match a+4, so it descends and finds a.
-    ret = graph_rewrite((a + 4) + (b + 5), _substitute, {a:c, b:c}, bottom_up=True, walk=True)
-    self.assertIs(ret, (c + 4) + (c + 5))
-
-  # *** bidirectional walk: bpm fires before children, pm fires after rebuild ***
-
-  def test_walk_bidirectional_visit_order(self):
-    """Bidirectional walk: bpm fires pre-order, pm fires post-order."""
-    visited = []
-    def bpm_visit(ctx, x):
-      ctx.append((x.arg if x.op is Ops.CONST else x.op, "bpm"))
-      return None
-    def pm_visit(ctx, x):
-      ctx.append((x.arg if x.op is Ops.CONST else x.op, "pm"))
-      return None
-    bpm = PatternMatcher([(UPat(GroupOp.All, name="x"), bpm_visit)])
-    pm = PatternMatcher([(UPat(GroupOp.All, name="x"), pm_visit)])
-    a = UOp.const(dtypes.int, 1)
-    b = UOp.const(dtypes.int, 2)
-    graph_rewrite(a + b, pm, ctx=visited, bpm=bpm, walk=True)
-    # bpm fires pre-order, pm fires post-order
-    self.assertEqual(visited, [
-      (Ops.ADD, "bpm"), (1, "bpm"), (1, "pm"), (2, "bpm"), (2, "pm"), (Ops.ADD, "pm"),
-    ])
-
-  def test_walk_bidirectional_bpm_short_circuits(self):
-    """If bpm matches, children are skipped and pm never fires on that node."""
-    visited = []
-    def bpm_match(ctx, x):
-      ctx.append((x.arg if x.op is Ops.CONST else x.op, "bpm"))
-      # rewrite const(1) -> const(10), short-circuiting its subtree
-      if x.op is Ops.CONST and x.arg == 1: return x.replace(arg=10)
-      return None
-    def pm_match(ctx, x):
-      ctx.append((x.arg if x.op is Ops.CONST else x.op, "pm"))
-      return None
-    bpm = PatternMatcher([(UPat(GroupOp.All, name="x"), bpm_match)])
-    pm = PatternMatcher([(UPat(GroupOp.All, name="x"), pm_match)])
-    a = UOp.const(dtypes.int, 1)
-    b = UOp.const(dtypes.int, 2)
-    ret = graph_rewrite(a + b, pm, ctx=visited, bpm=bpm, walk=True)
-    # bpm matches const(1) and short-circuits it, so pm never fires on const(1)
-    self.assertNotIn((1, "pm"), visited)
-    # but pm still fires on const(2) and the rebuilt ADD
-    self.assertIn((2, "pm"), visited)
-    self.assertIs(ret, UOp.const(dtypes.int, 10) + b)
-
 if __name__ == '__main__':
  unittest.main()
@@ -8,7 +8,7 @@ class TestDataset(unittest.TestCase):
    X_train[0].contiguous().realize()
    GlobalCounters.reset()
    X_train[0].contiguous().realize()
-    self.assertLessEqual(GlobalCounters.kernel_count, 1)  # 0 if BUFFER_VIEW (zero-copy), 1 otherwise
+    self.assertEqual(GlobalCounters.kernel_count, 1)

 if __name__ == '__main__':
  unittest.main()
@@ -5,7 +5,6 @@ class TestMultiRamUsage(unittest.TestCase):
  def setUp(self):
    gc.collect()
    self.baseline = GlobalCounters.mem_used
-    self.baseline_per_device = dict(GlobalCounters.mem_used_per_device)
    self.N = 100
  def assertUsed(self, amt, strict=True):
    gc.collect()
@@ -13,11 +12,6 @@ class TestMultiRamUsage(unittest.TestCase):
    print(f"used {used} bytes")
    if strict: self.assertEqual(used, amt)
    else: self.assertLessEqual(used, amt)
-  def assertDeviceUsed(self, expected:dict[str, int]):
-    gc.collect()
-    for dev, amt in expected.items():
-      used = GlobalCounters.mem_used_per_device[dev] - self.baseline_per_device.get(dev, 0)
-      self.assertEqual(used, amt, f"device {dev}: expected {amt} bytes used, got {used}")

  def test_zeros(self):
    _ = Tensor.zeros(self.N, self.N).contiguous().realize()
@@ -65,33 +59,6 @@ class TestMultiRamUsage(unittest.TestCase):
    X.shard_(devices_4, axis=0).realize()
    self.assertUsed(256 * 4)  # TODO: can be zero

-  def test_zeros_per_device(self):
-    _ = Tensor.zeros(self.N, self.N, device="NULL").contiguous().realize()
-    self.assertDeviceUsed({"NULL": self.N*self.N*4})
-
-  def test_zeros_del_per_device(self):
-    _ = Tensor.zeros(self.N, self.N, device="NULL").contiguous().realize()
-    del _
-    self.assertDeviceUsed({"NULL": 0})
-
-  def test_zeros_copy_per_device(self):
-    devices_2 = ("NULL:1", "NULL:2")
-    _ = Tensor.zeros(self.N, self.N).contiguous().to(devices_2).realize()
-    self.assertDeviceUsed({"NULL:1": self.N*self.N*4, "NULL:2": self.N*self.N*4})
-
-  def test_zeros_shard_per_device(self):
-    devices_2 = ("NULL:1", "NULL:2")
-    _ = Tensor.zeros(self.N, self.N).contiguous().shard(devices_2, axis=0).realize()
-    self.assertDeviceUsed({"NULL:1": self.N*(self.N//2)*4, "NULL:2": self.N*(self.N//2)*4})
-
-  def test_sharded_memory_replicated_per_device(self):
-    devices_4 = tuple(f"NULL:{i+1}" for i in range(4))
-    X = Tensor.ones(256, device="NULL").contiguous().realize()
-    self.assertDeviceUsed({"NULL": 256*4})
-    X.shard_(devices_4).realize()
-    for d in devices_4:
-      self.assertDeviceUsed({d: 256*4})
-
  def _test_matmul_half(self, dev_count:int):
    N = 32
    total_mem = {}
@@ -117,7 +117,7 @@ class TestContiguous(unittest.TestCase):
  def test_size_change_buffer_view(self):
    a = Tensor.empty(4)
    b = a.reshape((1, 1, 4)).shrink(((0, 1), (0, 1), (0, 3))).contiguous()
-    check_schedule(b, 0)  # contiguous shrink of a realized buffer is a zero-copy BUFFER_VIEW
+    check_schedule(b, 1)

  def test_double_contiguous_realizes_once(self):
    a = Tensor.empty(4, 1)
@@ -234,18 +234,6 @@ class TestSchedule(unittest.TestCase):
    d = Tensor.empty(1).assign(c)
    check_schedule(d, 1)

-  def test_detach_assign(self):
-    a = Tensor.ones(4, 4).contiguous().realize()
-    buf1, buf2 = Tensor.empty(4, 4).contiguous(), Tensor.empty(4, 4).contiguous()
-    r = buf2.assign(buf1.assign(a + 1.0) * 2.0)
-    check_schedule(r.detach().contiguous(), 2)
-
-  def test_contiguous_backward_assign(self):
-    a = Tensor.ones(4, 4).contiguous().realize()
-    buf1, buf2 = Tensor.empty(4, 4).contiguous(), Tensor.empty(4, 4).contiguous()
-    r = buf2.assign(buf1.assign(a + 1.0) * 2.0)
-    check_schedule(r.contiguous_backward().contiguous(), 2)
-
  def test_mulacc_relu_fusion(self):
    a = Tensor.empty(10)
    b = Tensor.empty(10)
@@ -1116,7 +1104,6 @@ class TestUOpBecome(unittest.TestCase):
    from tinygrad.helpers import all_same
    assert all_same([x.uop.base.realized for x in [a,b,c]])

-  @unittest.skip("not clear if we want this")
  def test_setitem_becomes_subbuffer(self):
    a = Tensor.full((4,), 2.).contiguous().realize()
    b = a.shrink(((0, 2),)).assign(Tensor.full((2,), 1.0))
@@ -1170,23 +1157,5 @@ class TestFusionOp(unittest.TestCase):
    self.assertEqual(len(sched), 1)
    self.assertLess(time.perf_counter()-st, 2.0)

-# NOTE: the NULL backend supports BUFFER_VIEW
-class TestBufferView(unittest.TestCase):
-  def test_shrink_contiguous_is_buffer_view(self):
-    # 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))
-
-  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))
-
-  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))
-
 if __name__ == '__main__':
  unittest.main(verbosity=2)
@@ -1,4 +1,4 @@
-import unittest, decimal, sys
+import unittest, decimal, json, struct, sys
 from dataclasses import dataclass
 from typing import Generator

@@ -282,10 +282,9 @@ class TestVizIntegration(BaseTestViz):
    ast = Tensor.schedule(Tensor.empty(4)+Tensor.empty(4))[0].ast
    prg = get_program(ast, Device[Device.DEFAULT].renderer)
    lst = get_viz_list()
-    self.assertEqual(len(lst), 3)
-    self.assertEqual(lst[0]["name"], "Process 1 Buffer n1")
-    self.assertEqual(lst[1]["name"], "Schedule 1 Kernel n1")
-    self.assertEqual(lst[2]["name"], prg.name)
+    self.assertEqual(len(lst), 2)
+    self.assertEqual(lst[0]["name"], "Schedule 1 Kernel n1")
+    self.assertEqual(lst[1]["name"], prg.name)

  # schedule graph CALL nodes have a link to jump to codegen
  def test_link_sched_codegen(self):
@@ -294,9 +293,8 @@ class TestVizIntegration(BaseTestViz):
    sched = Tensor.schedule(c1, c2)
    prgs = [si.lower().prg.p.name for si in sched]
    lst = get_viz_list()
-    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")
-    graph = next(get_viz_details(sched_idx, viz_kernel))["graph"]
+    viz_kernel = next(i for i,s in enumerate(lst[0]["steps"]) if s["name"] == "View Kernel Graph")
+    graph = next(get_viz_details(0, viz_kernel))["graph"]
    call_nodes = [n for n in graph.values() if n["label"].startswith("CALL")]
    for i,n in enumerate(call_nodes):
      assert n["ref"] is not None
@@ -357,9 +355,41 @@ class TestVizIntegration(BaseTestViz):

 from tinygrad.device import ProfileDeviceEvent, ProfileGraphEvent, ProfileGraphEntry
 from tinygrad.viz.serve import get_profile
-from extra.viz.cli import decode_profile

-def load_profile(lst:list[ProfileEvent]) -> dict: return decode_profile(get_profile(lst))
+class TinyUnpacker:
+  def __init__(self, buf): self.buf, self.offset = buf, 0
+  def __call__(self, fmt:str) -> tuple:
+    ret = struct.unpack_from(fmt, self.buf, self.offset)
+    self.offset += struct.calcsize(fmt)
+    return ret
+
+# 0 means None, otherwise it's an enum value
+def option(i:int) -> int|None: return None if i == 0 else i-1
+
+def load_profile(lst:list[ProfileEvent]) -> dict:
+  ret = get_profile(lst)
+  u = TinyUnpacker(ret)
+  total_dur, global_peak, index_len, layout_len = u("<IQII")
+  strings, dtypes, markers = json.loads(ret[u.offset:u.offset+index_len]).values()
+  u.offset += index_len
+  layout:dict[str, dict] = {}
+  for _ in range(layout_len):
+    klen = u("<B")[0]
+    k = ret[u.offset:u.offset+klen].decode()
+    u.offset += klen
+    layout[k] = v = {"events":[]}
+    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["peak"] = u("<Q")[0]
+      for _ in range(event_count):
+        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}

 class TestVizProfiler(BaseTestViz):
  def test_transfer_uses_copy_device(self):
@@ -29,7 +29,6 @@ class TestAssign(unittest.TestCase):
    a.realize()
    np.testing.assert_allclose(b.numpy(), 0)

-  @unittest.skip("TODO: this often crashes in CI")
  def test_assign_zeros(self):
    a = Tensor.zeros(10,10).contiguous()
    b = Tensor.zeros(10,10).contiguous()
@@ -270,16 +269,6 @@ class TestAssign(unittest.TestCase):
    out = attn.cache_k.flatten().numpy()
    np.testing.assert_allclose(out, [1.,1.,1.,1.,1.,1.,0.,0.,1.,1.,1.,1.,1.,1.,0.,0.])

-  def test_assign_after(self):
-    t = Tensor.zeros(10).contiguous().realize()
-    t.uop = t.uop.after(t.uop.assign((t+1).uop))
-    np.testing.assert_allclose(t.numpy(), [1.,1.,1.,1.,1.,1.,1.,1.,1.,1.])
-
-  def test_assign_after_partial(self):
-    t = Tensor.zeros(10).contiguous().realize()
-    t.uop = t.uop.after(t[:5].uop.assign(Tensor.ones(5).uop))
-    np.testing.assert_allclose(t.numpy(), [1.,1.,1.,1.,1.,0.,0.,0.,0.,0.])
-
  def test_assign_contiguous(self):
    b = Tensor.arange(16).reshape(4,4).contiguous().realize()
    a = (Tensor.arange(16).reshape(4,4).contiguous().realize() + 1)
@@ -496,10 +485,10 @@ class TestAssign(unittest.TestCase):
    np.testing.assert_allclose(c.numpy(), [4.0, 3.0, 3.0, 4.0])

  def test_assign_bitcast_different_size(self):
-    # assign to a shape-changing bitcast view (only works on DISK currently)
+    # different-size bitcast creates a new tensor, not a view, so assign doesn't modify the original
    a = Tensor([0]*8, dtype=dtypes.uint8).realize()
    a.bitcast(dtypes.int64).assign(Tensor([12345], dtype=dtypes.int64)).realize()
-    np.testing.assert_equal(a.numpy(), [0]*8)  # TODO: should be [57, 48, 0, 0, 0, 0, 0, 0] (little-endian 12345)
+    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):
@@ -609,8 +598,8 @@ class TestAssign(unittest.TestCase):
      x = q + caches[i][:1]             # next layer also references the same CONTIGUOUS through q
    GlobalCounters.reset()
    caches[-1][:1].contiguous().realize()
-    # 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)
+    # 2 kernels for first assign + 3 per remaining assign (matmul, contiguous, assign) + 1 final read = 3*N
+    self.assertEqual(GlobalCounters.kernel_count, 3*N)


 class TestAssignOrdering(unittest.TestCase):
@@ -687,16 +676,16 @@ class TestAssignOrdering(unittest.TestCase):
    """Swap two non-overlapping slices - requires reading both before writing."""
    # without .realize() on temps: values not captured before overwriting
    buf = Tensor([1, 2, 3, 4, 5, 6, 7, 8]).contiguous().realize()
-    left = buf[0:4].clone()  # lazy - not captured yet
-    right = buf[4:8].clone()  # lazy - not captured yet
+    left = buf[0:4].contiguous()  # lazy - not captured yet
+    right = buf[4:8].contiguous()  # lazy - not captured yet
    buf[0:4].assign(right).realize()  # this works
    buf[4:8].assign(left).realize()  # left now reads from modified buf!
    np.testing.assert_equal(buf.numpy(), [5, 6, 7, 8, 5, 6, 7, 8])  # TODO: wrong! should be [5,6,7,8,1,2,3,4]

    # with .realize() on temps: values captured before writes
    buf = Tensor([1, 2, 3, 4, 5, 6, 7, 8]).contiguous().realize()
-    left = buf[0:4].clone().realize()
-    right = buf[4:8].clone().realize()
+    left = buf[0:4].contiguous().realize()
+    right = buf[4:8].contiguous().realize()
    buf[0:4].assign(right).realize()
    buf[4:8].assign(left).realize()
    np.testing.assert_equal(buf.numpy(), [5, 6, 7, 8, 1, 2, 3, 4])
@@ -767,12 +756,13 @@ class TestAssignOrdering(unittest.TestCase):
    np.testing.assert_equal(b.numpy(), [1, 2, 3, 4])

  def test_variable_slice_ordering(self):
-    """Variable-indexed slices - conflicting variable binds in same schedule are rejected."""
+    """Variable-indexed slices - tests symbolic dependency tracking."""
    v_i = Variable("i", 0, 3)
    buf = Tensor.zeros(4, 4).contiguous().realize()
    buf[v_i.bind(0):v_i.bind(0)+1, :].assign(Tensor.ones(1, 4))
    buf[v_i.bind(1):v_i.bind(1)+1, :].assign(Tensor.ones(1, 4) * 2)
-    with self.assertRaises(RuntimeError): buf[0:1, :].sum().item()
+    self.assertEqual(buf[0:1, :].sum().item(), 4)
+    self.assertEqual(buf[1:2, :].sum().item(), 8)

  def test_multi_step_assign_read_write_same_buffer(self):
    """Assign to m and param reading b, then update b, across multiple steps.
@@ -1,6 +1,6 @@
 import unittest
 import numpy as np
-from tinygrad import Tensor, function
+from tinygrad import Tensor
 from tinygrad.dtype import dtypes
 from tinygrad.uop.ops import UOp

@@ -100,43 +100,5 @@ class TestCall(unittest.TestCase):
    c = Tensor.call(a, b, fxn=a.as_param(0) + b.as_param(1))
    np.testing.assert_equal(c.numpy(), 2 * np.ones((10, 10)))

-class TestCallSchedule(unittest.TestCase):
-  def test_reshape_precompile(self):
-    a = Tensor.empty(4, 8).realize()
-    a = a.reshape(4,4,2).assign(Tensor.empty(4,4,2)).reshape(8,4)
-    @function(precompile=True)
-    def s(x): return x.sum(axis=0)
-    (s(a)*3).realize()
-
-  def test_call_precompiled(self):
-    a = Tensor.empty(4, 8)
-    @function(precompile=True)
-    def s(x): return x*2
-    (s(a)*3).realize()
-
-  def test_double_call(self):
-    a = Tensor.empty(4, 8)
-    @function(precompile=True)
-    def s(x): return x*2
-    s(s(a)).realize()
-
-  def test_double_call_contiguous(self):
-    a = Tensor.empty(4, 8)
-    @function(precompile=True)
-    def s(x): return x*2
-    s(s(a).contiguous()).realize()
-
-  def test_call_double_gemm(self):
-    a = Tensor.randn(4, 8, requires_grad=True)
-    b = Tensor.randn(8, 12, requires_grad=True)
-    c = Tensor.randn(12, 16, requires_grad=True)
-    ref = Tensor.randn(4, 16)
-    Tensor.realize(a,b,c,ref)
-    @function(precompile=True)
-    def gemm(a:Tensor, b:Tensor, c:Tensor) -> Tensor: return (a@b)@c
-    out = gemm(a,b,c)
-    (out-ref).square().mean().backward()
-    out.realize(a.grad, b.grad, c.grad)
-
 if __name__ == '__main__':
  unittest.main()
@@ -77,6 +77,15 @@ class TestCallify(unittest.TestCase):
    out.callify()
    self.assertListEqual(out.tolist(), [5, 7, 9])

+  def test_callify_then_schedule(self):
+    a = Tensor([1.,2,3])
+    b = Tensor([4.,5,6])
+    out = a + b
+    out.callify()
+    schedule = out.schedule()
+    self.assertGreater(len(schedule), 0)
+    self.assertListEqual(out.tolist(), [5.0, 7.0, 9.0])
+
  def test_reduce(self):
    out = Tensor([1.,2,3,4]).sum()
    out.callify()
@@ -74,13 +74,13 @@ class TestRawDiskBuffer(unittest.TestCase):
    _test_bitcasted(t, dtypes.float32, 0.0)
    _test_bitcasted(t, dtypes.uint32, 0)
    # pi in float16 stored via int16
-    t.bitcast(dtypes.uint16).assign(Tensor.full((128, 64), 0x4248, dtype=dtypes.uint16)).realize()
+    t.assign(Tensor.full((128, 64), 0x4248, dtype=dtypes.uint16).bitcast(dtypes.uint8)).realize()
    _test_bitcasted(t, dtypes.float16, 3.140625)
    _test_bitcasted(t, dtypes.float32, 50.064727)
    _test_bitcasted(t, dtypes.uint16, 0x4248)
    _test_bitcasted(t, dtypes.uint32, 0x42484248)
    # pi in float32 stored via float32
-    t.bitcast(dtypes.float32).assign(Tensor.full((128, 32), 3.1415927, dtype=dtypes.float32)).realize()
+    t.assign(Tensor.full((128, 32), 3.1415927, dtype=dtypes.float32).bitcast(dtypes.uint8)).realize()
    _test_bitcasted(t, dtypes.float32, 3.1415927)
    _test_bitcasted(t, dtypes.uint32, 0x40490FDB)
    # doesn't suport normal cast
@@ -178,13 +178,6 @@ class TestSafetensors(TempDirTestCase):
    import json
    assert json.loads(dat[8:8+sz])['__metadata__']['hello'] == 'world'

-  def test_safe_save_only_copy(self):
-    from tinygrad.helpers import GlobalCounters
-    t = Tensor.rand(10, 10).realize()
-    GlobalCounters.reset()
-    safe_save({"t": t}, self.tmp("test_copy.safetensors"))
-    assert GlobalCounters.global_ops == 0, f"safe_save should have no compute, got {GlobalCounters.global_ops} ops"
-
  def test_save_all_dtypes(self):
    for dtype in dedup(DTYPES_DICT.values()):
      if dtype in [dtypes.bfloat16]: continue # not supported in numpy
@@ -364,10 +357,15 @@ class TestDiskTensor(TempDirTestCase):

  def test_assign_with_bitcast(self):
    # bitcast assign is used in safe_save for writing header length
+    # bitcast on source side works, bitcast on target side raises
    t = Tensor.empty(16, device=f"disk:{self.tmp('dt_assign_bitcast')}", dtype=dtypes.uint8)
-    t[0:8].bitcast(dtypes.int64).assign([12345])
+    # correct way: bitcast the source to match target dtype
+    t[0:8].assign(Tensor([12345], dtype=dtypes.int64, device="CPU").bitcast(dtypes.uint8))
    val = int.from_bytes(t[0:8].data(), 'little')
    self.assertEqual(val, 12345)
+    # bitcast on target with non-broadcastable dtype raises
+    with self.assertRaises(RuntimeError):
+      t[0:4].bitcast(dtypes.int32).assign(Tensor([12345], dtype=dtypes.int64))

  def test_assign_to_bitcast_view(self):
    # assign float values to a float32 view of a uint8 disk buffer (used by safe_save)
@@ -1,339 +0,0 @@
-import numpy as np
-import unittest
-from tinygrad.function import function
-from tinygrad import Tensor
-from tinygrad.uop.ops import UOp
-
-class TestFunction(unittest.TestCase):
-  def test_simple(self):
-    @function
-    def f(a:Tensor, b:Tensor) -> Tensor: return a+b
-
-    a = Tensor([1,2,3])
-    b = Tensor([4,5,6])
-    np.testing.assert_equal(f(a,b).numpy(), [5,7,9])
-
-  def test_simple_same(self):
-    @function
-    def f(a:Tensor, b:Tensor) -> Tensor: return a+b
-
-    a = Tensor([1,2,3])
-    np.testing.assert_equal(f(a,a).numpy(), [2,4,6])
-
-  def test_implicit(self):
-    inp = Tensor([7,8,9])
-    @function
-    def f(a:Tensor, b:Tensor) -> Tensor: return a+b+inp
-
-    a = Tensor([1,2,3])
-    b = Tensor([4,5,6])
-    np.testing.assert_equal(f(a,b).numpy(), [12,15,18])
-
-  def test_implicit_same_as_input(self):
-    inp = Tensor([7,8,9])
-    @function
-    def f(a:Tensor, b:Tensor) -> Tensor: return a+b+inp
-
-    a = Tensor([1,2,3])
-    np.testing.assert_equal(f(a, inp).numpy(), [15,18,21])
-
-  def test_implicit_2(self):
-    inp = Tensor([7,8,9])
-    @function
-    def f(a:Tensor, b:Tensor) -> Tensor:
-      return a+b+inp
-    inp2 = Tensor([7,8,10])
-    @function
-    def g(a:Tensor, b:Tensor) -> Tensor:
-      return a+b+inp2
-
-    a = Tensor([1,2,3])
-    b = Tensor([4,5,6])
-    c = f(a,b)
-    d = g(a,b)
-    c.realize(d)
-    np.testing.assert_equal(c.numpy(), [12,15,18])
-    np.testing.assert_equal(d.numpy(), [12,15,19])
-
-  def test_implicit_unrealized(self):
-    inp = Tensor([1,2,3]) + Tensor([4,5,6])
-    @function
-    def f(a:Tensor) -> Tensor: return a + inp
-
-    np.testing.assert_equal(f(Tensor([10,20,30])).numpy(), [15,27,39])
-
-  def test_detach(self):
-    @function
-    def f(a:Tensor, b:Tensor) -> Tensor: return a.detach() + b
-
-    a = Tensor([1,2,3])
-    b = Tensor([4,5,6])
-    np.testing.assert_equal(f(a, b).numpy(), [5,7,9])
-
-  def test_contiguous_backward(self):
-    @function
-    def f(a:Tensor, b:Tensor) -> Tensor: return (a + b).contiguous_backward()
-
-    a = Tensor([1,2,3])
-    b = Tensor([4,5,6])
-    np.testing.assert_equal(f(a, b).numpy(), [5,7,9])
-
-  def test_method(self):
-    class Foo:
-      def __init__(self): self.w = Tensor([10,20,30])
-      @function
-      def __call__(self, x:Tensor) -> Tensor: return x + self.w
-
-    foo = Foo()
-    np.testing.assert_equal(foo(Tensor([1,2,3])).numpy(), [11,22,33])
-
-  def test_grad_gemm(self):
-    @function
-    def f(a:Tensor, b:Tensor) -> Tensor: return a @ b
-
-    a = Tensor([[1.,2.],[3.,4.]], requires_grad=True)
-    b = Tensor([[5.,6.],[7.,8.]], requires_grad=True)
-    (f(a, b).contiguous() * b).sum().backward()
-    Tensor.realize(a, b, a.grad, b.grad)
-    # L = sum((a@b) * b), dL/d(a@b) = b, dL/da = b @ b^T, dL/db = a^T @ b + (a@b)
-    na, nb = a.numpy(), b.numpy()
-    np.testing.assert_allclose(a.grad.numpy(), nb @ nb.T)
-    np.testing.assert_allclose(b.grad.numpy(), na.T @ nb + na @ nb)
-
-  def test_grad_implicit(self):
-    w = Tensor([1., 2., 3.], requires_grad=True)
-    w.realize() # TODO: this is required
-    @function
-    def f(x:Tensor) -> Tensor: return x * w
-
-    x = Tensor([4., 5., 6.])
-    f(x).sum().backward()
-    np.testing.assert_allclose(w.grad.numpy(), [4., 5., 6.])
-
-  def test_symbolic_index(self):
-    table = Tensor([10,20,30,40]).contiguous().realize()
-    @function
-    def f(x:Tensor, start_pos:int|UOp) -> Tensor:
-      return x + table[start_pos]
-
-    v = UOp.variable("start_pos", 0, 3)
-    np.testing.assert_equal(f(Tensor([1,2,3]), v.bind(0)).numpy(), [11,12,13])
-
-  def test_symbolic_shape_input(self):
-    table = Tensor([10,20,30,40]).contiguous().realize()
-    @function
-    def f(x:Tensor) -> Tensor: return x * 2
-    sz = UOp.variable("sz", 1, 3)
-    slic = table[:sz.bind(2)]
-    np.testing.assert_equal(f(slic)[:2].numpy(), [20,40])
-
-  def test_nested_calls(self):
-    w = Tensor([10., 20., 30.])
-    @function
-    def f(a:Tensor) -> Tensor: return a + w
-    @function
-    def g(a:Tensor) -> Tensor: return a * w
-
-    a = Tensor([1., 2., 3.])
-    np.testing.assert_allclose(g(f(a)).numpy(), [110., 440., 990.])
-
-  def test_nested_calls_backward(self):
-    w = Tensor([[1., 2.], [3., 4.]]).contiguous().realize()
-    @function
-    def inner(x:Tensor) -> Tensor: return x + w
-    @function
-    def outer(a:Tensor, b:Tensor) -> Tensor: return inner(a.reshape(1,2) + b.reshape(1,2))
-
-    a = Tensor([1., 2.], requires_grad=True)
-    b = Tensor([3., 4.], requires_grad=True)
-    outer(a, b).sum().backward()
-    np.testing.assert_allclose(a.grad.numpy(), [2., 2.])
-    np.testing.assert_allclose(b.grad.numpy(), [2., 2.])
-
-  def test_unused_param_backward(self):
-    @function
-    def f(a:Tensor, b:Tensor, c:Tensor) -> Tensor: return a + c  # b is unused
-
-    a = Tensor([1., 2., 3.], requires_grad=True)
-    b = Tensor([4., 5., 6.], requires_grad=True)
-    c = Tensor([7., 8., 9.], requires_grad=True)
-    f(a, b, c).sum().backward()
-    np.testing.assert_allclose(a.grad.numpy(), [1., 1., 1.])
-    np.testing.assert_allclose(b.grad.numpy(), [0., 0., 0.])
-    np.testing.assert_allclose(c.grad.numpy(), [1., 1., 1.])
-
-  def test_name(self):
-    @function
-    def f(a:Tensor) -> Tensor: return a + 1
-    assert f(Tensor([1])).uop.arg.name.endswith("f")
-
-  def test_method_name(self):
-    class Foo:
-      @function
-      def __call__(self, x:Tensor) -> Tensor: return x + 1
-    assert Foo()(Tensor([1])).uop.arg.name.endswith("Foo.__call__")
-
-  def test_callable_instance(self):
-    class Foo:
-      def __init__(self): self.w = Tensor([10,20,30])
-      def __call__(self, x:Tensor) -> Tensor: return x + self.w
-    foo = Foo()
-    f = function(foo)
-    np.testing.assert_equal(f(Tensor([1,2,3])).numpy(), [11,22,33])
-    assert f(Tensor([1,2,3])).uop.arg.name.endswith("Foo")
-
-  def test_iadd(self):
-    @function
-    def f(x:Tensor) -> Tensor:
-      x += 1
-      return x
-
-    a = Tensor([1,2,3]).realize()
-    np.testing.assert_equal(f(a).numpy(), [2,3,4])
-    np.testing.assert_equal(a.numpy(), [3,4,5])  # TODO: should be [1,2,3]
-
-  def test_implicit_assign(self):
-    a = Tensor([1,2,3])
-    a += 1
-    c = Tensor([2,2,2]).contiguous()
-    @function
-    def f(b:Tensor) -> Tensor: return a+b+c
-    b = Tensor([10,20,30]).realize()
-    np.testing.assert_equal(f(b).numpy(), [14,25,36])
-
-  def test_assign_input(self):
-    @function
-    def f(a:Tensor, b:Tensor) -> Tensor:
-      a.assign(b+1)
-      return a
-
-    a = Tensor([1,2,3]).realize()
-    b = Tensor([10,20,30]).realize()
-    np.testing.assert_equal(f(a,b).numpy(), [11,21,31])
-    np.testing.assert_equal(a.numpy(), [11,21,31])  # TODO: should be [1,2,3]
-    np.testing.assert_equal(b.numpy(), [10,20,30])
-
-  @unittest.expectedFailure
-  def test_assign_slice(self):
-    @function
-    def f(a:Tensor, b:Tensor) -> Tensor:
-      a[1:] = b[1:]+1
-      return a
-
-    a = Tensor([1,2,3]).realize()
-    b = Tensor([10,20,30]).realize()
-    np.testing.assert_equal(f(a,b).numpy(), [1,21,31])
-    np.testing.assert_equal(a.numpy(), [1,2,3])
-    np.testing.assert_equal(b.numpy(), [10,20,30])
-
-class TestFunctionMulti(unittest.TestCase):
-  devices_2 = ("CPU:0", "CPU:1")
-
-  def test_simple_multi(self):
-    @function
-    def f(a:Tensor, b:Tensor) -> Tensor: return a+b
-
-    a = Tensor([1,2,3,4]).shard(self.devices_2, axis=None)
-    b = Tensor([10,20,30,40]).shard(self.devices_2, axis=None)
-    np.testing.assert_equal(f(a,b).numpy(), [11,22,33,44])
-
-  def test_simple_multi_sharded(self):
-    @function
-    def f(a:Tensor, b:Tensor) -> Tensor: return a+b
-
-    a = Tensor([1,2,3,4]).shard(self.devices_2, axis=0)
-    b = Tensor([10,20,30,40]).shard(self.devices_2, axis=0)
-    np.testing.assert_equal(f(a,b).numpy(), [11,22,33,44])
-
-  def test_data_parallel_multi(self):
-    @function
-    def f(x:Tensor, w:Tensor) -> Tensor: return x @ w
-
-    x = Tensor([[1.,2.],[3.,4.],[5.,6.],[7.,8.]]).shard(self.devices_2, axis=0)
-    w = Tensor([[1.,0.],[0.,1.]]).shard(self.devices_2, axis=None)
-    np.testing.assert_allclose(f(x, w).numpy(), [[1.,2.],[3.,4.],[5.,6.],[7.,8.]])
-
-  def test_grad_implicit_multi(self):
-    w = Tensor([1., 2., 3., 4.], requires_grad=True).shard(self.devices_2, axis=None)
-    w.realize()
-    @function
-    def f(x:Tensor) -> Tensor: return x * w
-
-    x = Tensor([4., 5., 6., 7.]).shard(self.devices_2, axis=None)
-    f(x).sum().backward()
-    np.testing.assert_allclose(w.grad.numpy(), [4., 5., 6., 7.])
-
-  def test_call_axis(self):
-    @function
-    def f(x:Tensor, w:Tensor) -> Tensor: return x @ w
-
-    x = Tensor([[1.,0.],[0.,1.],[1.,1.],[0.,0.]]).shard(self.devices_2, axis=0)
-    w = Tensor([[1.,2.],[3.,4.]]).shard(self.devices_2, axis=None)
-    result = f(x, w)
-    # CALL output should inherit axis=0 from the sharded input
-    self.assertEqual(result.uop.axis, 0)
-    # reduce on the sharded axis should remove it
-    self.assertIsNone(result.sum().uop.axis)
-
-  def test_call_axis_shard_inside(self):
-    @function
-    def f(x:Tensor, w:Tensor) -> Tensor:
-      return x.shard(self.devices_2, axis=0) @ w.shard(self.devices_2, axis=None)
-
-    x = Tensor([[1.,0.],[0.,1.],[1.,1.],[0.,0.]])
-    w = Tensor([[1.,2.],[3.,4.]])
-    result = f(x, w)
-    self.assertEqual(result.uop.axis, 0)
-    np.testing.assert_allclose(result.numpy(), x.numpy() @ w.numpy())
-
-  def test_data_parallel_backward(self):
-    @function
-    def f(x:Tensor, w:Tensor) -> Tensor: return x @ w
-
-    x = Tensor([[1.,0.],[0.,1.],[1.,1.],[0.,0.]], requires_grad=True).shard(self.devices_2, axis=0)
-    w = Tensor([[1.,2.],[3.,4.]], requires_grad=True).shard(self.devices_2, axis=None)
-    w.realize()
-    f(x, w).sum().backward()
-    # d/dx = ones @ w^T = [[1,3],[1,3],[1,3],[1,3]], but sum so ones(4,2) @ w^T? no:
-    # L = sum(x @ w), dL/dx = ones(4,2) @ w^T... actually dL/d(xw) = ones(4,2), dL/dx = ones(4,2) @ w^T
-    np.testing.assert_allclose(x.grad.numpy(), np.ones((4,2)) @ np.array([[1,3],[2,4]]))
-
-  def test_data_parallel_backward_4(self):
-    devices_4 = tuple(f"CPU:{i}" for i in range(4))
-    @function
-    def f(x:Tensor, w:Tensor) -> Tensor: return x @ w
-
-    x = Tensor(np.arange(16).reshape(8,2).astype(np.float32), requires_grad=True).shard(devices_4, axis=0)
-    w = Tensor([[1.,2.],[3.,4.]], requires_grad=True).shard(devices_4, axis=None)
-    w.realize()
-    f(x, w).sum().backward()
-    np.testing.assert_allclose(x.grad.numpy(), np.ones((8,2)) @ np.array([[1,3],[2,4]]))
-
-  def test_data_parallel_backward_implicit(self):
-    devices_4 = tuple(f"CPU:{i}" for i in range(4))
-    w = Tensor([[1.,2.],[3.,4.]], requires_grad=True).shard(devices_4, axis=None)
-    w.realize()
-    @function
-    def f(x:Tensor) -> Tensor: return x @ w
-
-    x = Tensor(np.arange(16).reshape(8,2).astype(np.float32), requires_grad=True).shard(devices_4, axis=0)
-    f(x).sum().backward()
-    np.testing.assert_allclose(x.grad.numpy(), np.ones((8,2)) @ np.array([[1,3],[2,4]]))
-
-  def test_data_parallel_backward_twice(self):
-    devices_4 = tuple(f"CPU:{i}" for i in range(4))
-    w = Tensor([[1.,2.],[3.,4.]], requires_grad=True).shard(devices_4, axis=None)
-    w.realize()
-    # pre-init grads like the training loop does
-    w.grad = w.zeros_like().contiguous().realize()
-    @function
-    def f(x:Tensor) -> Tensor: return x @ w
-
-    expected = np.ones((8,2)) @ np.array([[1,3],[2,4]])
-    for _ in range(2):
-      x = Tensor(np.arange(16).reshape(8,2).astype(np.float32), requires_grad=True).shard(devices_4, axis=0)
-      f(x).sum().backward()
-      np.testing.assert_allclose(x.grad.numpy(), expected)
-
-if __name__ == '__main__':
-  unittest.main()
@@ -1,4 +1,4 @@
-import os, struct, unittest
+import os, unittest
 from tinygrad import dtypes, Tensor, fetch, Device
 from tinygrad.nn.state import ggml_data_to_tensor, gguf_load
 from tinygrad.device import is_dtype_supported
@@ -120,41 +120,5 @@ class TestGGUF(unittest.TestCase):
      else:
        self.assertEqual(kv_data[k], read_val(-1))

-class TestGGUFGEMV(unittest.TestCase):
-  def _test_gguf_gemv(self, qtype: GGMLQuantizationType):
-    block_size, type_size = GGML_QUANT_SIZES[qtype]
-    rows, cols = 8192, 2048
-    n_blocks = rows * cols // block_size
-    rng = np.random.default_rng(42)
-    # 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 == GGMLQuantizationType.Q8_0: q_data[:, :2] = scales[:, :2]                    # d at offset 0
-    elif qtype == GGMLQuantizationType.Q4_K: q_data[:, :4] = scales[:, :4]                  # d, dmin at offset 0
-    elif qtype == GGMLQuantizationType.Q6_K: q_data[:, -2:] = scales[:, :2]                 # d at end
-    q_data = q_data.flatten()
-    ref = dequantize(q_data, qtype).reshape(rows, cols)
-
-    # build a minimal gguf in memory: header + 1 tensor info + aligned data
-    buf = bytearray()
-    buf += struct.pack("<4siqq", b"GGUF", 3, 1, 0)              # magic, version, n_tensors, n_kv
-    buf += struct.pack("<Q", 6) + b"weight"                      # tensor name
-    buf += struct.pack("<I", 2)                                  # ndims
-    buf += struct.pack("<QQ", cols, rows)                        # dims (gguf stores reversed)
-    buf += struct.pack("<i", qtype.value)
-    buf += struct.pack("<Q", 0)                                  # offset
-    buf += b"\x00" * ((32 - len(buf) % 32) % 32)                # pad to alignment=32
-    buf += q_data.tobytes()
-
-    _, tensors = gguf_load(Tensor(np.frombuffer(buf, dtype=np.uint8)).to(None))
-
-    x = rng.standard_normal(cols).astype(np.float32)
-    np.testing.assert_allclose((tensors["weight"] @ Tensor(x)).numpy(), ref @ x, atol=1e-2, rtol=1e-2)
-    np.testing.assert_equal(tensors["weight"].numpy(), ref)
-
-  def test_gguf_gemv_q8_0(self): self._test_gguf_gemv(GGMLQuantizationType.Q8_0)
-  def test_gguf_gemv_q4_k(self): self._test_gguf_gemv(GGMLQuantizationType.Q4_K)
-  def test_gguf_gemv_q6_k(self): self._test_gguf_gemv(GGMLQuantizationType.Q6_K)
-
 if __name__ == '__main__':
  unittest.main()
@@ -68,14 +68,6 @@ class TestTensorGradient(unittest.TestCase):
    np.testing.assert_allclose(x.grad.numpy(), [2.0+3.0+2*3.0])
    self.assertIs(x.grad, old_grad)

-  def test_gradient_through_chained_unrealized_setitem(self):
-    g1 = Tensor.zeros(4).contiguous()
-    g1[2] = Tensor(1.0)
-    g2 = Tensor.zeros(5, 4).contiguous()
-    g2[0] = g1
-    x = Tensor.randn(4, 4)
-    np.testing.assert_allclose(x.pad(((1,0),(0,0))).gradient(x, gradient=g2)[0].numpy(), np.zeros((4, 4)))
-
 class TestViewGradient(unittest.TestCase):
  def test_expand(self):
    x = Tensor.randn(5,2)
@@ -179,6 +179,8 @@ class TestIndexing(unittest.TestCase):
    def delitem(): del reference[0]
    self.assertRaises(TypeError, delitem)

+  # TODO setitem backward
+  '''
  def test_set_item_to_scalar_tensor(self):
    m = random.randint(1, 10)
    n = random.randint(1, 10)
@@ -188,6 +190,7 @@ class TestIndexing(unittest.TestCase):
    z[:, 0] = w
    z.sum().backward()
    numpy_testing_assert_equal_helper(w.grad, m * a)
+  '''

  def test_step(self):
    v = Tensor.arange(10)
@@ -83,15 +83,6 @@ class TestLinAlg(unittest.TestCase):
      s_diag = (S.unsqueeze(-2) * Tensor.eye(2))
      reconstruction_helper([U, s_diag, V], a)

-  def test_svd_identity_4x4(self):
-    a = Tensor.eye(4)
-    U,S,V = a.svd()
-    assert not np.isnan(U.numpy()).any()
-    assert not np.isnan(S.numpy()).any()
-    assert not np.isnan(V.numpy()).any()
-    s_diag = (S.unsqueeze(-2) * Tensor.eye(4))
-    reconstruction_helper([U, s_diag, V], a)
-
  def test_svd_rank1(self):
    a = Tensor([[1.0, 1.0], [2.0, 2.0]]).realize()
    U, S, V = a.svd()
@@ -45,31 +45,37 @@ class TestTinyFS(unittest.TestCase):
    cls._server.shutdown()
    cls._server.server_close()

+  @unittest.expectedFailure
  def test_store(self):
    h = Tensor([1.0, 2.0, 3.0, 4.0]).fs_store().realize()
    self.assertEqual(h.shape, (16,))
    self.assertEqual(h.dtype, dtypes.uint8)

+  @unittest.expectedFailure
  def test_store_deterministic(self):
    a = Tensor([1.0, 2.0, 3.0, 4.0]).fs_store().realize()
    b = Tensor([1.0, 2.0, 3.0, 4.0]).fs_store().realize()
    np.testing.assert_array_equal(a.numpy(), b.numpy())

+  @unittest.expectedFailure
  def test_store_different_data(self):
    a = Tensor([1.0, 2.0, 3.0, 4.0]).fs_store().realize()
    b = Tensor([5.0, 6.0, 7.0, 8.0]).fs_store().realize()
    self.assertNotEqual(a.tolist(), b.tolist())

+  @unittest.expectedFailure
  def test_roundtrip_uint8(self):
    arr = np.arange(256, dtype=np.uint8)
-    loaded = Tensor(arr).fs_store().realize().fs_load(len(arr)).to("CPU")
+    loaded = Tensor(arr).fs_store().realize().fs_load(len(arr))
    np.testing.assert_array_equal(loaded.numpy(), arr)

+  @unittest.expectedFailure
  def test_roundtrip_multichunk_uint8(self):
    arr = np.random.default_rng(42).integers(0, 256, size=Tensor.CHUNK_SIZE + 1024, dtype=np.uint8)
-    loaded = Tensor(arr).fs_store().realize().fs_load(len(arr)).to("CPU")
+    loaded = Tensor(arr).fs_store().realize().fs_load(len(arr))
    np.testing.assert_array_equal(loaded.numpy(), arr)

+  @unittest.expectedFailure
  def test_hash_matches_python_impl(self):
    arr = np.arange(256, dtype=np.uint8)
    h = Tensor(arr).fs_store().realize()
@@ -4,7 +4,6 @@ if int(os.getenv("TYPED", "0")):
  install_import_hook(__name__)
 from tinygrad.tensor import Tensor                                    # noqa: F401
 from tinygrad.engine.jit import TinyJit                               # noqa: F401
-from tinygrad.function import function                                # noqa: F401
 from tinygrad.uop.ops import UOp
 Variable = UOp.variable
 from tinygrad.dtype import dtypes                                     # noqa: F401
@@ -1,6 +1,6 @@
 from __future__ import annotations
 import sys, argparse, typing, re, unicodedata, json, uuid, time, functools
-from tinygrad import Tensor, nn, UOp, TinyJit, getenv, function
+from tinygrad import Tensor, nn, UOp, TinyJit, getenv
 from tinygrad.helpers import partition, DEBUG, Timing, GlobalCounters, stderr_log, colored
 from tinygrad.viz.serve import TCPServerWithReuse, HTTPRequestHandler

@@ -116,7 +116,6 @@ class TransformerBlock:
      self.ffn_up      = nn.Linear(dim, hidden_dim, bias=False)
      self.ffn_down    = nn.Linear(hidden_dim, dim, bias=False)

-  @function
  def _attention(self, x:Tensor, start_pos:int|UOp) -> Tensor:
    x_norm = self.attn_norm(x)                       # (B,T,D)
    q, k, v = self.attn_q(x_norm), self.attn_k(x_norm), self.attn_v(x_norm)
@@ -132,25 +131,19 @@ class TransformerBlock:
    q = apply_rope(q, freqs_cis)
    k = apply_rope(k, freqs_cis)

-    # TODO: fix assign to behave like this
-    assigned_kv = self.cache_kv.uop.after(self.cache_kv[:, :, :, start_pos:start_pos+T, :].uop.assign(Tensor.stack(k, v).contiguous().uop))
-    tensor_assigned_kv = Tensor(assigned_kv, device=assigned_kv.device)
-    k = tensor_assigned_kv[0, :, :, 0:start_pos+T, :]
-    v = tensor_assigned_kv[1, :, :, 0:start_pos+T, :]
-
-    #self.cache_kv[:, :, :, start_pos:start_pos+T, :].assign(Tensor.stack(k, v))
-    #k = self.cache_kv[0, :, :, 0:start_pos+T, :]
-    #v = self.cache_kv[1, :, :, 0:start_pos+T, :]
+    if not hasattr(self, "cache_kv"):
+      self.cache_kv = Tensor.zeros(2, B, self.n_kv_heads, self.max_context, self.head_dim, dtype=k.dtype, device=k.device).contiguous().realize()
+    self.cache_kv[:, :, :, start_pos:start_pos+T, :].assign(Tensor.stack(k, v))
+    k = self.cache_kv[0, :, :, 0:start_pos+T, :]
+    v = self.cache_kv[1, :, :, 0:start_pos+T, :]

    # NOTE: this mask is causal_lower_right, not the causal_upper_left generated by is_casual = True
-    # TODO: this if statement should be removed and it shouldn't generate extra kernels
-    mask = Tensor.full((1, 1, T, start_pos+T), float("-inf"), dtype=x.dtype, device=x.device).triu(start_pos+1) if T > 1 else None
+    mask = Tensor.full((1, 1, T, start_pos+T), float("-inf"), dtype=x.dtype, device=x.device).triu(int(start_pos)+1) if T > 1 else None
    attn = q.scaled_dot_product_attention(k, v, attn_mask=mask, enable_gqa=True)     # (B,H,T,Hd)
    attn = attn.transpose(1, 2).reshape(B, T, -1)                                    # back to (B,T,D)
    attn = self.attn_output(attn)
    return x + attn

-  @function(precompile=bool(getenv("PRECOMPILE", 0)))
  def _feed_forward(self, h: Tensor) -> Tensor:
    h_norm = self.ffn_norm(h)
    if hasattr(self, 'ffn_gate_exps'):
@@ -163,9 +156,6 @@ class TransformerBlock:
    return h + self.ffn_down(gated)

  def __call__(self, x: Tensor, start_pos: int|UOp):
-    if not hasattr(self, "cache_kv"):
-      # TODO: how is the dtype of this determined?
-      self.cache_kv = Tensor.zeros(2, x.shape[0], self.n_kv_heads, self.max_context, self.head_dim, device=x.device).contiguous().realize()
    return self._feed_forward(self._attention(x, start_pos)).contiguous()

 class Transformer:
@@ -190,9 +180,9 @@ class Transformer:
    return (self.forward_jit if getenv("JIT", 1) and tokens.shape[1] == 1 and isinstance(start_pos, UOp) else self.forward)(tokens, start_pos)

  @staticmethod
-  def from_gguf(gguf:Tensor, max_context:int|None=None, realize=bool(getenv("REALIZE", 1))) -> tuple[Transformer, dict]:
+  def from_gguf(gguf:Tensor, max_context:int|None=None, realize=True) -> tuple[Transformer, dict]:
    # TODO: remove the need for copy to default device
-    kv, state_dict = nn.state.gguf_load(gguf.to(None).realize())
+    kv, state_dict = nn.state.gguf_load(gguf.to(None))

    # all state items should be float16, not float32
    state_dict = {k:v.cast('float16') if getenv("HALF", 1) else v for k,v in state_dict.items()}
@@ -220,9 +210,8 @@ class Transformer:
                        num_experts=kv.get(f'{arch}.expert_count', 0), num_experts_per_tok=kv.get(f'{arch}.expert_used_count', 0))
    nn.state.load_state_dict(model, state_dict, verbose=False, consume=True, realize=False)  # NOTE: rope_freqs.weight (32,) is unused
    # NOTE: without this contiguous, it unpacks the weights from the model every time. we shouldn't need this, but for now it's faster
-    if realize:
-      for s in (params:=nn.state.get_parameters(model)): s.replace(s.contiguous())
-      Tensor.realize(*params)
+    for s in (params:=nn.state.get_parameters(model)): s.replace(s.contiguous())
+    if realize: Tensor.realize(*params)
    return model, kv

  def generate(self, tokens:list[int], start_pos=0):
@@ -264,7 +253,7 @@ CHAT_HTML = b'''<!DOCTYPE html><html><head><title>tinygrad chat</title><style>
 </style></head><body><div id="chat"></div>
 <textarea id="input" rows="1" placeholder="Ask anything"></textarea>
 <script>
-  input.onkeydown = (e) => { if (e.key === 'Enter' && !e.shiftKey && !e.isComposing) { e.preventDefault(); send() } }
+  input.onkeydown = (e) => { if (e.key === 'Enter' && !e.shiftKey) { e.preventDefault(); send() } }
  const msgs = [];
  async function send() {
    if (!input.value.trim()) return;
@@ -338,42 +327,37 @@ class Handler(HTTPRequestHandler):

 if __name__ == "__main__":
  parser = argparse.ArgumentParser()
-  parser.add_argument("--model", "-m", choices=list(models.keys()), default=list(models.keys())[0], help="Model choice")
+  parser.add_argument("--model", choices=list(models.keys()), default=list(models.keys())[0], help="Model choice")
  parser.add_argument("--max_context", type=int, default=4096, help="Max Context Length")
  parser.add_argument("--serve", nargs='?', type=int, const=11434, metavar="PORT", help="Run OpenAI compatible API (optional port, default 11434)")
  parser.add_argument("--benchmark", nargs='?', type=int, const=20, metavar="COUNT", help="Benchmark tok/s (optional count, default 20)")
  args = parser.parse_args()

  # load the model
-  raw_model = Tensor.from_url(models[args.model])
-  model, kv = Transformer.from_gguf(raw_model, args.max_context)
-  if DEBUG >= 1 or args.benchmark:
-    print(f"using model {args.model} with {raw_model.nbytes():,} bytes and {sum(x.numel() for x in nn.state.get_parameters(model)):,} params")
-  del raw_model
+  model, kv = Transformer.from_gguf(Tensor.from_url(models[args.model]), args.max_context)
+  if DEBUG >= 1: print(f"using model {args.model}")

-  # TODO: why this is required to free the RAM of the GGUF copy?
-  import gc
-  gc.collect()
+  # do benchmark
+  if args.benchmark:
+    param_bytes = sum(x.nbytes() for x in nn.state.get_parameters(model))
+    for b in model.blk:
+      if hasattr(b, 'ffn_gate_exps'):
+        expert_bytes = b.ffn_gate_exps.weight.nbytes() + b.ffn_up_exps.weight.nbytes() + b.ffn_down_exps.weight.nbytes()
+        param_bytes -= int(expert_bytes * (1 - b.num_experts_per_tok / b.ffn_gate_exps.weight.shape[0]))
+    gen = model.generate([0], 0)
+    for _ in range(args.benchmark):
+      GlobalCounters.reset()
+      with Timing(on_exit=lambda x: f", {1e9/x:6.2f} tok/s, {GlobalCounters.global_mem/x:7.2f} GB/s, param {param_bytes/x:7.2f} GB/s"): next(gen)
+    exit(0)

  # extract some metadata
  tok = SimpleTokenizer.from_gguf_kv(kv)
  bos_id: int|None = kv.get('tokenizer.ggml.bos_token_id') if kv.get('tokenizer.ggml.add_bos_token', True) else None
  eos_id: int = kv['tokenizer.ggml.eos_token_id']

-  # do benchmark
-  if args.benchmark:
-    gen = model.generate(toks:=[bos_id or 0], 0)
-    for _ in range(args.benchmark):
-      GlobalCounters.reset()
-      with Timing(on_exit=lambda x: f", {1e9/x:6.2f} tok/s, {GlobalCounters.global_mem/x:7.2f} GB/s,"
-                  f" {GlobalCounters.global_mem//1000000}/{GlobalCounters.mem_used//1000000} MB  --  "+\
-                  tok.decode(toks).replace("\n", "\\n")): next(gen)
-    exit(0)
-
  # start server
  if args.serve: TCPServerWithReuse(('', args.serve), Handler).serve_forever()

-  # interactive chat
  ids: list[int] = [bos_id] if bos_id is not None else []
  while 1:
    start_pos = max(len(ids) - 1, 0)
@@ -41,15 +41,15 @@ def full_rewrite_to_sink(sink:UOp, ren:Renderer|None=None, optimize:bool=True) -
    # split ranges
    sink = graph_rewrite(sink, pm_split_ranges+pm_flatten_range, ctx={}, name="split ranges")

-    # create image buffers
-    if IMAGE == 1 and ren.device in {"QCOM", "CL"}: sink = graph_rewrite(sink, pm_make_images, name="create image buffers", bottom_up=True)
-
    # symbolic (NOTE: this is a requirement for pm_simplify_ranges to be correct)
    sink = graph_rewrite(sink, sym+pm_flatten_range, name="initial symbolic")

    # optimize (schedule) the AST
    sink = graph_rewrite(sink, pm_simplify_ranges, name="simplify ranges")

+    # create image buffers
+    if IMAGE == 1 and ren.device in {"QCOM", "CL"}: sink = graph_rewrite(sink, pm_make_images, name="create image buffers", bottom_up=True)
+
    # do postrange optimization, BEAM or hand_coded_optimizations
    sink = apply_opts(sink, ren)

@@ -1,7 +1,7 @@
 from typing import Any, cast
 import functools, itertools
 from collections import defaultdict
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from tinygrad.dtype import dtypes, ImageDType, DType, AddrSpace, Invalid, PtrDType
 from tinygrad.uop.ops import UOp, Ops, UPat, PatternMatcher, GroupOp, identity_element
 from tinygrad.uop.symbolic import uop_given_valid, parse_valid, invalid_gate
@@ -308,6 +308,8 @@ pm_render = PatternMatcher([
@dataclass
 class ReduceContext:
  acc_num: int = 0
+  # track ENDs by range for merging parallel reduces
+  range_to_ends: dict[tuple[UOp, ...], list[UOp]] = field(default_factory=dict)

 def horizontal_reduce(inp:UOp, out_dtype:DType) -> list[UOp]:
  # if this has a horizontal reduction component, do that first
@@ -333,15 +335,13 @@ def reduce_to_acc(ctx:ReduceContext, red:UOp):
    ctx.acc_num += 1
  ret = functools.reduce(lambda x,y: x.alu(red.arg, y), lst)
  if len(reduce_range) == 0: return ret
-  end = acc.index(UOp.const(dtypes.int, 0)).store(ret).end(*reduce_range).rtag("mergeable")
+  end = acc.index(UOp.const(dtypes.int, 0)).store(ret).end(*reduce_range)
+  ctx.range_to_ends.setdefault(reduce_range, []).append(end)
  return acc.after(end).index(UOp.const(dtypes.int, 0))

 def merge_reduce_ends(ctx:ReduceContext, sink:UOp):
-  # merge ENDs that share the same range (only those created by reduce_to_acc)
-  range_to_ends: dict[tuple[UOp, ...], list[UOp]] = {}
-  for u in sink.backward_slice:
-    if u.op is Ops.END and u.tag == "mergeable": range_to_ends.setdefault(u.src[1:], []).append(u)
-  subs = {e: UOp.group(*(e.src[0] for e in ends)).end(*r) for r, ends in range_to_ends.items() if len(ends) > 1 for e in ends}
+  # merge ENDs that share the same range
+  subs = {e: UOp.group(*(e.src[0] for e in ends)).end(*r) for r, ends in ctx.range_to_ends.items() if len(ends) > 1 for e in ends}
  return sink.substitute(subs) if subs else None

 pm_reduce = PatternMatcher([
@@ -36,8 +36,7 @@ def get_test_global_size(global_size, max_global_size, var_vals):
  return test_global_size, input_size / prod(test_global_size)

 def _time_program(p:ProgramSpec, lib:bytes, var_vals:dict[str, int], rawbufs:list[Buffer], early_stop:float|None=None,
-                  allow_test_size:int=True, max_global_size:int|None=65536, clear_l2=False, cnt=3, name="test", dev_timeout=False) -> list[float]:
-  timeout = int(early_stop * 1e3) if dev_timeout and early_stop is not None and early_stop < math.inf else None
+                  allow_test_size:int=True, max_global_size:int|None=65536, clear_l2=False, cnt=3, name="test") -> list[float]:
  factor = 1
  if allow_test_size and max_global_size is not None:
    global_size, factor = get_test_global_size(p.global_size, max_global_size, var_vals)
@@ -51,7 +50,7 @@ def _time_program(p:ProgramSpec, lib:bytes, var_vals:dict[str, int], rawbufs:lis
      if hasattr(dev:=Device[p.device], 'invalidate_caches'): dev.invalidate_caches()
      else:
        with Context(DEBUG=0, BEAM=0, CAPTURING=0, TRACK_MATCH_STATS=0): Tensor.ones(1024,1024).contiguous().realize(do_update_stats=False)
-    tms.append(unwrap(car(input_bufs, var_vals, wait=True, timeout=timeout))*factor)
+    tms.append(unwrap(car(input_bufs, var_vals, wait=True))*factor)
    if early_stop is not None and early_stop < min(tms): break
  return tms

@@ -162,8 +161,7 @@ def beam_search(s:Scheduler, rawbufs:list[Buffer], amt:int, allow_test_size=True
          continue
        seen_libs.add(lib)
        try: tms = _time_program(p, lib, var_vals, rawbufs, early_stop=beam[0][1]*3 if len(beam) else 1.0,
-                                 allow_test_size=allow_test_size, clear_l2=hasattr(dev, 'invalidate_caches'),
-                                 dev_timeout=getenv("BEAM_DEV_TIMEOUT", 1))
+                                 allow_test_size=allow_test_size, clear_l2=hasattr(dev, 'invalidate_caches'))
        except Exception as e:
          if BEAM_DEBUG: print(f"BEAM failed for opts: {candidates[i].applied_opts}\n{e}")
          if isinstance(e, RuntimeError): continue
@@ -141,7 +141,6 @@ class Buffer:
      self._buf = opaque if opaque is not None else self.allocator.alloc(self.nbytes, self.options)
      if not self.device.startswith("DISK") and (self.options is None or self.options.external_ptr is None):
        GlobalCounters.mem_used += self.nbytes
-        GlobalCounters.mem_used_per_device[self.device] += self.nbytes
      if PROFILE: Buffer.profile_events.append(ProfilePointEvent(self.device, "alloc", self.trace_num, {"dtype":self.dtype, "sz":self.size}))
    return self
  def deallocate(self):
@@ -150,7 +149,6 @@ class Buffer:
    if self._base is None:
      if GlobalCounters is not None and not self.device.startswith("DISK") and (self.options is None or self.options.external_ptr is None):
        GlobalCounters.mem_used -= self.nbytes
-        GlobalCounters.mem_used_per_device[self.device] -= self.nbytes
      if PROFILE: Buffer.profile_events.append(ProfilePointEvent(self.device, "free", self.trace_num))
      self.allocator.free(self._buf, self.nbytes, self.options)
    elif self._base is not None: self._base.allocated_views -= 1
@@ -1,7 +1,7 @@
 from dataclasses import dataclass, field
-from tinygrad.uop.ops import UOp, UPat, PatternMatcher, Ops, GroupOp, graph_rewrite, track_rewrites
-from tinygrad.dtype import dtypes, ImageDType
-from tinygrad.helpers import prod, DEBUG, argsort, VIZ, pluralize, FLOAT16
+from tinygrad.uop.ops import UOp, UPat, PatternMatcher, Ops, GroupOp, graph_rewrite, identity_element, profile_matches
+from tinygrad.dtype import ImageDType
+from tinygrad.helpers import prod, DEBUG, argsort, VIZ

@dataclass
 class AllocCtx:
@@ -18,29 +18,31 @@ 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"))
+  to_disk = isinstance(u._device, str) and u._device.startswith("DISK")
  if to_disk: ctx.buffer_map[u] = UOp.new_buffer(u.device, u.shard_size, u.dtype).reshape(u.max_shard_shape)
  # 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"])
+  from_creation = isinstance(u.src[0]._device, str) and any(u.src[0]._device.startswith(x) for x in ["NPY", "DISK", "PYTHON"])
  if from_creation: return tag_uop(ctx, u)

 def apply_after(ctx:AllocCtx, u:UOp):
-  base = u.src[0]
-  while base.op is Ops.AFTER: base = base.src[0]
-  ctx.buffer_map[u] = base
+  ctx.buffer_map[u] = u.src[0]

 # CONTIGUOUS and ASSIGN + parents are the only nodes that get updated
 add_tags = PatternMatcher([
  (UPat(Ops.COPY, name="u"), disk_copy_is_buffer),
-  # no tag on copies/allreduces that are assigned
-  (UPat(Ops.ASSIGN, src=(UPat(), UPat((Ops.COPY, Ops.ALLREDUCE), name="c")), name="a"),
+  # no tag on copies that are assigned
+  (UPat(Ops.ASSIGN, src=(UPat(), UPat(Ops.COPY, name="c")), name="a"),
   lambda a,c: a.replace(src=(a.src[0], c.rtag(())), tag=a.tag+c.tag) if a.tag and c.tag else None),
  (UPat(Ops.AFTER, name="u"), apply_after),
  (UPat({Ops.CONTIGUOUS, Ops.ASSIGN}, name="x"), tag_uop),
  (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:
+def replace_contig_with_assign(u:UOp):
+  # if size is 0, remove the contig
+  if u.size == 0: return u.src[0]
+  # no real contig for DISK tensors, they are left alone
+  if isinstance(u._device, str) and u._device.startswith("DISK"): return u.rtag(None)
  dtype = u.dtype
  if isinstance(dtype, ImageDType):
    if prod(dtype.shape) != prod(u.max_shard_shape) or ([x for x in u.max_shard_shape if x != 1] or [1])[-1] % 4 != 0:
@@ -48,86 +50,45 @@ def _buffer_like(u:UOp) -> UOp:
      dtype = dtype.base
  buffer = UOp.new_buffer(u.device, u.shard_size, dtype).reshape(u.max_shard_shape)
  if isinstance(u.device, tuple) and u.axis is not None: buffer = buffer.multi(u.axis)
-  return buffer
-
-def replace_contig_with_assign(u:UOp):
-  # if size is 0, remove the contig
-  if u.size == 0: 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)
-  return _buffer_like(u).assign(u.src[0]).rtag(u.tag)
+  return buffer.assign(u.src[0]).rtag(u.tag)

 def replace_assign_with_contig(u:UOp):
  assigned_to = u
-  while assigned_to.op in {Ops.ASSIGN, Ops.BITCAST, Ops.AFTER}: assigned_to = assigned_to.src[0].base
+  while assigned_to.op in {Ops.ASSIGN, Ops.BITCAST}: assigned_to = assigned_to.src[0].base
  if assigned_to.op is not Ops.BUFFER:
    return u.src[1].contiguous(tag=u.tag)

 def found_contiguous(ctx:dict[UOp, UOp], contig:UOp, src:UOp):
-  if (x:=src).op is Ops.CAST and x.dtype == dtypes.half and FLOAT16: x, contig = x.src[0], contig.cast(dtypes.float)
-  while x is not x.base:
+  x = src
+  while x is not src.base:
    if x.op is Ops.PERMUTE: contig = contig.permute(argsort(x.marg))
    elif x.op is Ops.RESHAPE: contig = contig.reshape(x.src[0].shape)
    else: return None
    x = x.src[0]
-  ctx[x] = contig
+  ctx[src.base] = contig

-def contiguous_mops_to_view(c:UOp):
-  """CONTIGUOUS(MOPS(BUFFER)) → CONTIGUOUS(BUFFER_VIEW) when movement ops collapse to a contiguous range."""
-  src = c.src[0]
-  buf = src.base
-  if buf.op not in {Ops.BUFFER, Ops.BUFFER_VIEW}: return None
-  if src.op is Ops.RESHAPE and src.src[0].op in {Ops.BUFFER, Ops.BUFFER_VIEW}: return None
-
-  # no symbolic shape
-  if not all(isinstance(x, int) for x in c.shape): return None
-
-  # check if view is supported
-  if not isinstance(c.device, str): return None
-  from tinygrad.device import Device
-  if not hasattr(Device[c.device].allocator, "_offset"): return None
-
-  # see if this can be a view
-  offset = src.contiguous_view_offset()
-  if offset is None: return None
-
-  # merge BUFFER_VIEWs
-  if buf.op is Ops.BUFFER_VIEW: offset, buf = offset + buf.arg[1], buf.src[0]
-
-  # NOTE: this contiguous is removed because this BUFFER_VIEW/RESHAPE has_buffer_identity
-  return UOp(Ops.BUFFER_VIEW, src.dtype, (buf,), (src.size, offset)).reshape(src.shape).contiguous(tag=c.tag)
-
-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
-  out = _buffer_like(c)
-  fxn = out.param_like(len(c.src)-1).assign(c.src[0]).sink()
-  return out.after(c.replace(src=(fxn,)+tuple(x.contiguous() if x.op is not Ops.AFTER else x for x in c.src[1:])+(out,), dtype=dtypes.void, tag=None))
-
-# 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),
-
-  # CONTIGUOUS(MOPS(BUFFER/BUFFER_VIEW)) → CONTIGUOUS(BUFFER_VIEW) when movement ops collapse to contiguous range
-  (UPat(Ops.CONTIGUOUS, src=(UPat(GroupOp.Movement),), name="c"), contiguous_mops_to_view),
-
-  # *** CONTIGUOUS replacement hack for openpilot ***
-  (UPat(Ops.CONTIGUOUS, src=(UPat((*GroupOp.Movement, Ops.CAST), name="src"),), name="contig"), found_contiguous),
+  # CONTIGUOUS replacement hack for openpilot
+  (UPat(Ops.CONTIGUOUS, src=(UPat(GroupOp.Movement, name="src"),), name="contig"), found_contiguous),
  # replace ALU sources with contiguous versions found above
  (UPat(GroupOp.ALU, name="alu"), lambda ctx,alu: alu.replace(src=new_src) if (new_src:=tuple(ctx.get(s, s) for s in alu.src)) != alu.src else None),
-
  # add CONTIGUOUS to tagged UOps
  (UPat(GroupOp.All-{Ops.CONTIGUOUS, Ops.ASSIGN}, name="x"), lambda x: x.rtag(None).contiguous(tag=x.tag) if x.tag else x.replace(tag=None)),
-  # remove extra CONTIGUOUS on ASSIGN (only when assign target is contiguous)
-  (UPat(Ops.CONTIGUOUS, src=(UPat(Ops.ASSIGN, name="a"),), name="c"),
-   lambda a,c: a.replace(tag=(a.tag or ())+(c.tag or ())) if a.src[0].has_buffer_identity() else None),
+  # remove extra CONTIGUOUS on ASSIGN
+  (UPat(Ops.CONTIGUOUS, src=(UPat(Ops.ASSIGN, name="a"),), name="c"), lambda a,c: a.replace(tag=a.tag+c.tag)),
  # replace ASSIGN with CONTIGUOUS
  (UPat(Ops.ASSIGN, name="u"), replace_assign_with_contig),
  # replace CONTIGUOUS with ASSIGNs
  (UPat(Ops.CONTIGUOUS, name="u"), replace_contig_with_assign),
-  # remove DETACH/CONTIGUOUS_BACKWARD (allows more contiguous removal)
+  # remove DETACH/CONTIGUOUS_BACKWARD
  (UPat((Ops.DETACH, Ops.CONTIGUOUS_BACKWARD), name="x"), lambda x: x.src[0]),
+  # reduce of size 0 is the identity element
+  (UPat(Ops.REDUCE_AXIS, name="reduce", src=(UPat.var("x"),)),
+   lambda reduce,x: reduce.const_like(identity_element(reduce.arg[0], reduce.dtype)) if x.size == 0 and reduce.size != 0 else None),
+  # handle size 0
+  (UPat(GroupOp.All-{Ops.SINK}, name="x"), lambda x: x.const_like(0).rtag(x.tag) if x._shape is not None and x.size == 0 else None),
+  # early fixup const copy (TODO: is this wrong if there's a pad?)
+  (UPat(Ops.COPY, src=(UPat.var("s"), UPat()), name="c"), lambda c,s: c.const_like(ss.arg) if (ss:=s.base).op is Ops.CONST else None),
 ])

 def untag_and_append(ctx:AllocCtx, x:UOp):
@@ -152,7 +113,7 @@ def replace_input_buffer(ctx:AllocCtx, b:UOp):
 pm_finalize_call = PatternMatcher([
  (UPat(Ops.ASSIGN, 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.COPY, name="x"), lambda ctx,x: append_after(ctx,x) if isinstance(x.device, str) and x.device.startswith("DISK") else None),
  # replace UNIQUE with LUNIQUE for CONST cache key normalization
  (UPat(Ops.CONST, src=(UPat(Ops.UNIQUE), UPat(Ops.DEVICE, name="d")), name="b"), lambda b,d: b.replace(src=(d,))),
 ])
@@ -160,15 +121,12 @@ pm_finalize_call = PatternMatcher([
 pm_replace_buf = PatternMatcher([
  # replace BUFFER with PARAM for cache key normalization
  (UPat(Ops.BUFFER, src=(UPat(Ops.UNIQUE), UPat(Ops.DEVICE)), name="b"), replace_input_buffer),
-  # replace BUFFER_VIEW with PARAM. this rewrite is bottom up so BUFFERs we don't need won't be in the input
-  (UPat(Ops.BUFFER_VIEW, src=(UPat(Ops.BUFFER),), name="b"), replace_input_buffer),
  # strip value from BIND for cache key normalization, so different values hit same cache
  (UPat(Ops.BIND, src=(UPat(Ops.DEFINE_VAR), UPat(Ops.CONST)), name="b"), replace_input_buffer),
 ])

-@track_rewrites(lambda _,ret: f"Process {pluralize('Buffer', len(ret[1]))}")
+@profile_matches
 def transform_to_call(big_sink:UOp) -> tuple[UOp, dict[UOp, UOp]]:
-  if VIZ: graph_rewrite(big_sink, PatternMatcher([]), name="View Tensor Graph")
  # uop list is a list in the original_sink graph and we can map to the tags later
  # here we build buffer map
  dont_realize = {Ops.CONST, Ops.BUFFER, Ops.BIND, Ops.DEFINE_VAR, Ops.AFTER}
@@ -183,6 +141,6 @@ def transform_to_call(big_sink:UOp) -> tuple[UOp, dict[UOp, UOp]]:

  # here we construct the final buffer_map. this is everything that will go into the tensor map
  graph_rewrite(big_sink, pm_finalize_call, ctx=ctx, name="finalize call")
-  ret = graph_rewrite(UOp.sink(*ctx.assigns), pm_replace_buf, ctx=ctx, bottom_up=True, name="replace bufs").call(*ctx.replacements)
+  ret = graph_rewrite(UOp.sink(*ctx.assigns), pm_replace_buf, ctx=ctx, name="replace bufs").call(*ctx.replacements)
  if VIZ: graph_rewrite(ret, PatternMatcher([]), name="View Call")
  return ret, ctx.buffer_map
@@ -50,7 +50,7 @@ class CompiledRunner(Runner):

  def __reduce__(self): return self.__class__, (self.p,)

-  def __call__(self, rawbufs:list[Buffer], var_vals:dict[str, int]|None=None, wait=False, timeout:int|None=None) -> float|None:
+  def __call__(self, rawbufs:list[Buffer], var_vals:dict[str, int]|None=None, wait=False) -> float|None:
    if var_vals is None: var_vals = {}
    global_size, local_size = self.p.launch_dims(var_vals)
    if Device[self.p.device].renderer.has_local and local_size is None and all_int(self.p.global_size):
@@ -58,7 +58,7 @@ class CompiledRunner(Runner):
      global_size = [g//l if g%l == 0 else g/l for g,l in zip(global_size, local_size)]
      self.p = replace(self.p, global_size=global_size, local_size=local_size)
    return self._prg(*[x._buf for x in rawbufs], global_size=tuple(global_size), local_size=tuple(local_size) if local_size else None,
-                     vals=tuple(var_vals[k.expr] if k.expr not in self.p.runtimevars else None for k in self.p.vars), wait=wait, timeout=timeout)
+                     vals=tuple(var_vals[k.expr] if k.expr not in self.p.runtimevars else None for k in self.p.vars), wait=wait)

 class ViewOp(Runner):
  def __init__(self, buf:Buffer): super().__init__(colored(f"view {buf.nbytes:8d} @ {buf.offset:<10d}", "yellow"), buf.device)
@@ -1,10 +1,10 @@
 import time, inspect
 from typing import cast
 from collections import deque
-from tinygrad.uop.ops import UOp, Ops, buffers, UOpMetaClass, track_rewrites, graph_rewrite, gate_kernel_sink, KernelInfo
+from tinygrad.uop.ops import UOp, Ops, KernelInfo, buffers, UOpMetaClass, track_rewrites, graph_rewrite, gate_kernel_sink
 from tinygrad.uop.spec import type_verify, tensor_spec
 from tinygrad.device import Buffer, MultiBuffer
-from tinygrad.helpers import DEBUG, cpu_profile, TracingKey, SPEC, pluralize, SCACHE, BASEDIR, flatten
+from tinygrad.helpers import DEBUG, cpu_profile, TracingKey, SPEC, pluralize, SCACHE, BASEDIR
 from tinygrad.engine.realize import ExecItem

 # **** schedule linearizer
@@ -22,7 +22,7 @@ def create_schedule(sched_sink:UOp) -> UOp:
    for u in sched_sink.toposort(gate_kernel_sink):
      if u.op is not Ops.AFTER: continue
      k = u.src[1]
-      assert k.op in {Ops.CALL, Ops.END}, f"AFTER src[1] should be CALL or END, not {k.op}"
+      assert k.op in {Ops.CALL, Ops.END}, f"AFTER src[1] should be KERNEL or END, not {k.op}"
      in_degree.setdefault(k, 0)
      if k.op is Ops.END: assert k.src[0].op is Ops.CALL, f"END src[0] should be KERNEL, not {k.src[0].op}"
      # WAR deps from rangeify are stored in AFTER src[2:]
@@ -49,18 +49,64 @@ def create_schedule(sched_sink:UOp) -> UOp:
    linearized: list[UOp] = []
    while len(queue):
      rk = queue.popleft()
-      if rk.op is Ops.LINEAR:
-        linearized.extend(rk.src)
-      else:
-        k = rk.src[0] if rk.op is Ops.END else rk
-        assert k.op is Ops.CALL, f"unexpected op in queue: {k.op}"
-        buf_uops = tuple(_unwrap_src(s).buf_uop for s in k.src[1:] if s.op is not Ops.BIND)
-        linearized.append(k.src[0].call(*buf_uops, metadata=k.arg.metadata))
+      k = rk.src[0] if rk.op is Ops.END else rk
+      assert k.op is Ops.CALL, f"unexpected op in queue: {k.op}"
+      buf_uops = tuple(_unwrap_src(s).buf_uop for s in k.src[1:] if s.op is not Ops.BIND)
+      linearized.append(k.src[0].call(*buf_uops, metadata=k.arg.metadata))
      for x in children.get(rk, []):
        in_degree[x] -= 1
        if in_degree[x] == 0: queue.append(x)
  return UOp(Ops.LINEAR, src=tuple(linearized))

+from tinygrad.engine.memory import memory_planner
+from tinygrad.schedule.rangeify import get_kernel_graph
+from tinygrad.uop.ops import PatternMatcher, UPat
+
+def create_new_buffer(ctx:tuple[dict[UOp, UOp], tuple[UOp, ...]], b:UOp):
+  if (ret:=ctx[0].get(b, None)) is None: ctx[0][b] = ret = UOp.new_buffer(b.device, b.arg, b.dtype)
+  return ret
+
+pm_post_sched_cache = PatternMatcher([
+  # tag=True prevents re-matching after replacement (needed when PARAMs replace with PARAMs in nested callify)
+  (UPat(Ops.PARAM, name="x"), lambda ctx,x: ctx[1][x.arg].replace(tag=True) if x.tag is None else None),
+  # create new BUFFERs for LUNIQUE BUFFERs from rangeify
+  (UPat(Ops.BUFFER, src=(UPat(Ops.LUNIQUE), UPat(Ops.DEVICE)), name="b"), create_new_buffer),
+])
+
+schedule_cache: dict[bytes, UOp] = {}
+
+def _resolve_params(linear:UOp, params:tuple[UOp, ...]) -> UOp:
+  """Replace PARAMs in a LINEAR with the given params (BUFFERs or outer PARAMs), also handling LUNIQUE BUFFERs."""
+  from tinygrad.uop.ops import _remove_all_tags
+  linear = graph_rewrite(linear, pm_post_sched_cache, ctx=({}, params), name="params to buffers")
+  return graph_rewrite(linear, _remove_all_tags, name="remove tags")
+
+def rewrite_call_to_linear(ctx:list, call:UOp) -> UOp|None:
+  """Rewrite rule: CALL(SINK, *params) -> LINEAR(...) with caching. Only matches top-level CALLs from transform_to_call."""
+  function = call.src[0]
+  if function.op is not Ops.SINK or isinstance(function.arg, KernelInfo): return None
+  # recursively schedule any nested CALLs inside the function (from nested callify)
+  inner_start = len(ctx)
+  function = graph_rewrite(function, pm_schedule, ctx=ctx, name="schedule nested calls")
+  if not SCACHE or (linear:=schedule_cache.get(function.key, None)) is None:
+    if SPEC: type_verify(call.replace(src=(function,)+call.src[1:]), tensor_spec)
+    linear = create_schedule(get_kernel_graph(function))
+    if SCACHE: schedule_cache[function.key] = linear
+  # late apply params to buffers (tag=True prevents PARAM->PARAM cycles in nested callify)
+  linear = _resolve_params(linear, call.src[1:])
+  # resolve remaining PARAMs in inner LINEARs from nested CALLs using this call's params
+  for i in range(inner_start, len(ctx)):
+    inner_call, inner_linear = ctx[i]
+    ctx[i] = (inner_call, _resolve_params(inner_linear, call.src[1:]))
+  ctx.append((call, linear))
+  return linear
+
+pm_schedule = PatternMatcher([
+  (UPat(Ops.CALL, name="call"), rewrite_call_to_linear),
+  # strip AFTER(buf, LINEAR) -> buf after scheduling
+  (UPat(Ops.AFTER, src=(UPat(name="buf"), UPat(Ops.LINEAR))), lambda ctx,buf: buf),
+])
+
 def linear_to_schedule(linear:UOp) -> list[ExecItem]:
  """Convert a LINEAR UOp to a list of ExecItems."""
  schedule: list[ExecItem] = []
@@ -79,92 +125,44 @@ def linear_to_schedule(linear:UOp) -> list[ExecItem]:
      for j, bufs in enumerate(zip(*[x.bufs for x in cast(tuple[MultiBuffer, ...], ubufs)])):
        schedule.append(ExecItem(ast, list(bufs), metadata, {dnums[0].expr:j} if len(dnums) else {}))
    else:
-      schedule.append(ExecItem(ast, cast(list[Buffer|None], ubufs), metadata))
+      schedule.append(ExecItem(ast, list(ubufs), metadata))
  return schedule

-from tinygrad.engine.memory import memory_planner
-from tinygrad.schedule.rangeify import get_kernel_graph
-from tinygrad.uop.ops import PatternMatcher, UPat
-
-def create_new_buffer(ctx:tuple[dict[UOp, UOp], tuple[UOp, ...]], b:UOp):
-  if (ret:=ctx[0].get(b, None)) is None: ctx[0][b] = ret = UOp.new_buffer(b.device, b.arg, b.dtype)
-  return ret
-
-pm_post_sched_cache = PatternMatcher([
-  (UPat(Ops.PARAM, name="x"), lambda ctx,x: ctx[1][x.arg]),
-  # create new BUFFERs for LUNIQUE BUFFERs from rangeify
-  (UPat(Ops.BUFFER, src=(UPat(Ops.LUNIQUE), UPat(Ops.DEVICE)), name="b"), create_new_buffer),
-])
-
-pm_resolve_linear_call = PatternMatcher([
-  # call LINEAR is resolved here
-  (UPat(Ops.CALL, src=(UPat(Ops.LINEAR),), name="linear_call", allow_any_len=True), lambda linear_call:
-   graph_rewrite(linear_call.src[0], pm_post_sched_cache, ctx=({}, linear_call.src[1:]), walk=True, name="params to buffers")),
-  # LINEAR on LINEAR
-  (UPat(Ops.LINEAR, custom_early_reject={Ops.LINEAR}, name="x"),
-   lambda x: x.replace(src=tuple(flatten(x.src if x.op is Ops.LINEAR else (x,) for x in x.src)))),
-])
-
-schedule_cache: dict[bytes, UOp] = {}
-# ctx is just for DEBUG on inner
-def lower_sink_to_linear(function:UOp) -> UOp|None:
+# strip AFTER(buf, LINEAR) -> buf, used by _apply_map_to_tensors to clean up scope tensors after scheduling
+@track_rewrites(lambda _,ret: f"Schedule {pluralize('Kernel', len(ret[1]))}")
+def complete_create_schedule_with_vars(big_sink:UOp) -> tuple[list[UOp], list[ExecItem], dict[str, int]]:
  st = time.perf_counter()
-  if isinstance(function.arg, KernelInfo): return None
-  cache_key = function.key
-  if not SCACHE or (sc_ret:=schedule_cache.get(cache_key, None)) is None:
-    if SPEC: type_verify(function, tensor_spec)
-    # support recursive CALLs
-    linear = create_schedule(get_kernel_graph(function))
-    if SCACHE: schedule_cache[cache_key] = linear
-  else:
-    # schedule cache hit
-    linear = sc_ret
-  if (DEBUG >= 1 and len(linear.src) > 1) or DEBUG >= 3:
+
+  # rewrite CALLs to LINEARs and strip AFTERs
+  call_linear_pairs: list[tuple[UOp, UOp]] = []
+  graph_rewrite(big_sink, pm_schedule, ctx=call_linear_pairs, name="schedule calls")
+
+  # collect ExecItems from all LINEARs
+  schedule: list[ExecItem] = []
+  for _, linear in call_linear_pairs:
+    schedule.extend(linear_to_schedule(linear))
+
+  # get var_vals from CALL params
+  used_vars = set().union(*[{v.expr for v in si.src[0].variables()} for _, linear in call_linear_pairs for si in linear.src])
+  var_vals: dict[str, int] = {}
+  for call, _ in call_linear_pairs:
+    for b in call.src[1:]:
+      if b.op is Ops.BIND:
+        nm = b.src[0].expr
+        if nm not in used_vars: continue
+        val = b.src[1].arg
+        assert nm not in var_vals or var_vals[nm] == val, f"bind mismatch on {nm}, {var_vals[nm]} != {val}"
+        var_vals[nm] = val
+
+  with cpu_profile(TracingKey("memory planner")): schedule = memory_planner(schedule)
+
+  if (DEBUG >= 1 and len(schedule) > 1) or DEBUG >= 3:
    for frm in inspect.stack():
-      if frm.filename == "<string>": continue
      if frm.filename.startswith(str(BASEDIR / "apps")): break
      if not frm.filename.startswith(str(BASEDIR)) and not frm.filename.endswith("/contextlib.py"): break
    else:
      frm = None
-    print(f"scheduled {len(linear.src):5d} kernels in {(time.perf_counter()-st)*1000:8.2f} ms"+\
-          f" | {' cache hit' if SCACHE and sc_ret is not None else 'CACHE MISS'} {cache_key.hex()[:8]}"+\
+    print(f"scheduled {len(schedule):5d} kernels in {(time.perf_counter()-st)*1000:8.2f} ms"+\
          f" | {len(UOpMetaClass.ucache):7d} uops in cache"+("" if frm is None else f" | {frm.filename}:{frm.lineno}"))
-  return linear

-def soft_allreduce(c:UOp, a:UOp):
-  from tinygrad.schedule.multi import handle_allreduce
-  to = c.src[1].param_like(0)
-  src = c.src[2].param_like(1)
-  red = UOp(Ops.ALLREDUCE, dtype=a.arg, src=(src, a.src[1]), arg=a.arg)
-  return to.assign(handle_allreduce(src, red)).sink().call(*c.src[1:])
-
-pm_schedule = PatternMatcher([
-  (UPat(Ops.SINK, name="function"), lower_sink_to_linear),
-  # soft handler of allreduce
-  (UPat(Ops.CALL, src=(UPat(Ops.ALLREDUCE, name="a"),), allow_any_len=True, name="c"), soft_allreduce),
-])
-
-@track_rewrites(lambda _,ret: f"Schedule {pluralize('Kernel', len(ret[0]))}")
-def complete_create_schedule_with_vars(big_sink:UOp) -> tuple[list[ExecItem], dict[str, int]]:
-  # big_sink srcs are all the Tensors
-  linear_call = graph_rewrite(big_sink, pm_schedule, name="schedule to linear", enter_calls=True)
-
-  # this recursively resolves the linear_call and allocates buffers
-  linear = graph_rewrite(linear_call, pm_resolve_linear_call, name="resolve linear call")
-
-  # vars used in the schedule
-  used_vars = set().union(*[{v.expr for v in si.src[0].variables()} for si in linear.src])
-  # get var_vals
-  var_vals: dict[str, int] = {}
-  for b in big_sink.src[1:]:
-    if b.op is Ops.BIND:
-      nm = b.src[0].expr
-      if nm not in used_vars: continue
-      val = b.src[1].arg
-      if var_vals.get(nm, val) != val: raise RuntimeError(f"bind mismatch on {nm}, {var_vals[nm]} != {val}")
-      var_vals[nm] = val
-
-  # convert LINEAR to ExecItems
-  schedule: list[ExecItem] = linear_to_schedule(linear)
-  with cpu_profile(TracingKey("memory planner")): schedule = memory_planner(schedule)
-  return schedule, var_vals
+  return [call for call, _ in call_linear_pairs], schedule, var_vals
@@ -1,71 +0,0 @@
-import functools
-from typing import Generic, TypeVar, Callable, cast, overload
-from tinygrad.helpers import Context, dedup, getenv
-from tinygrad.uop.ops import UOp, Ops, graph_rewrite, PatternMatcher, UPat
-from tinygrad.tensor import Tensor
-
-def add_to_ctx(ctx, x:UOp):
-  ret = x.param_like(len(ctx))
-  ctx.append(x)
-  return ret
-
-pm_ctx = PatternMatcher([
-  (UPat((Ops.BUFFER, Ops.BIND), name="x"), add_to_ctx),
-  (UPat((Ops.ASSIGN, Ops.CONTIGUOUS), name="x"),
-   lambda ctx,x: add_to_ctx(ctx,x) if not x.op_in_backward_slice_with_self(Ops.PARAM) else None),
-])
-
-ReturnType = TypeVar('ReturnType')
-class _function(Generic[ReturnType]):
-  def __init__(self, fxn:Callable[..., ReturnType], *, precompile:bool=False):
-    self.fxn = fxn
-    self.precompile = precompile
-
-  def __get__(self, obj, objtype=None): return functools.partial(self.__call__, obj) if obj is not None else self
-
-  def __call__(self, *args, **kwargs) -> ReturnType:
-    input_uops: list[UOp] = [(t.uop if isinstance(t, Tensor) else t)
-                             for name,t in list(enumerate(args))+sorted(kwargs.items()) if isinstance(t, (Tensor, UOp))]
-
-    # use the base
-    #input_uops = [x.multibase for x in input_uops]
-
-    # deduplicate input_uops, keeping the first occurrence index for each unique uop
-    call_uops: list[UOp] = dedup(input_uops)
-
-    # disable realize/schedule while this is running
-    # run it and do surgery later
-    with Context(ALLOW_DEVICE_USAGE=getenv("DEVICE_IN_FUNCTION_BUG", 0)):
-      ret = self.fxn(*args, **kwargs)
-    assert isinstance(ret, Tensor), "only supports one tensor return for now"
-
-    # replace the known inputs with params (using deduplicated slots)
-    subs = {}
-    for i,x in enumerate(call_uops): subs[x] = x.param_like(i)
-    uret = ret.uop.substitute(subs)
-
-    # add contiguous to call_uops
-    #call_uops = [x.contiguous() for x in call_uops]
-
-    # the BUFFERs that are left are the implicit inputs
-    uret = graph_rewrite(uret, pm_ctx, call_uops, bottom_up=True, name="get_implicit_inputs")
-    name = getattr(self.fxn, '__qualname__', None) or type(self.fxn).__qualname__
-
-    # assign output
-    #pbuffer = uret.param_like(len(call_uops))
-    #assigned = pbuffer.assign(uret).sink()
-    #buffer = UOp.new_buffer(pbuffer.device, pbuffer.size, pbuffer.dtype).reshape(uret.shape)
-    #call = assigned.call(*call_uops, buffer, name=name)
-    #ret = buffer.after(call)
-
-    ret = uret.call(*call_uops, name=name, precompile=self.precompile)
-    return cast(ReturnType, Tensor(ret, device=ret.device))
-
-# overload signatures support both @function and @function(precompile=True) syntax
-@overload
-def function(fxn:Callable[..., ReturnType], *, precompile:bool=False) -> _function[ReturnType]: ...
-@overload
-def function(fxn:None=None, *, precompile:bool=False) -> Callable[[Callable[..., ReturnType]], _function[ReturnType]]: ...
-def function(fxn=None, *, precompile:bool=False):
-  if fxn is None: return lambda f: _function(f, precompile=precompile)
-  return _function(fxn, precompile=precompile)
@@ -13,21 +13,14 @@ def reduce_gradient(ctx:UOp, ret:UOp, op:Ops):
    return ((mask/broadcast_to_input(count)) * broadcast_to_input(ctx),)
  if op == Ops.MUL: return (broadcast_to_input(ctx * ret) / ret.src[0],)

-def call_gradient(ctx:UOp, k:UOp) -> tuple[UOp|None, ...]:
+def call_gradient(ctx:UOp, k:UOp):
  if k.arg.grad_fxn is not None: return (None,) + k.arg.grad_fxn(ctx, k)
  # auto-differentiate the function
  fxn, args = k.src[0], k.src[1:]
-  params = {x.arg:x for x in fxn.toposort(enter_calls=False) if x.op == Ops.PARAM}
-  grads = compute_gradient(fxn, ctx.param_like(len(args)), set(params.values()))
-  ret: list[UOp|None] = [None]
-  for i in range(len(args)):
-    if (p:=params.get(i, None)) is not None and p in grads:
-      # TODO: compact the args and remove unused ones
-      assert not grads[p].op_in_backward_slice_with_self(Ops.BUFFER), "BUG: BUFFER in backward slice of grad"
-      ret.append(grads[p].call(*args, ctx, name=(k.arg.name or "")+f"_backward_{i}"))
-    else:
-      ret.append(None)
-  return tuple(ret)
+  params = sorted([x for x in fxn.toposort() if x.op == Ops.PARAM], key=lambda x: x.arg)
+  grads = compute_gradient(fxn, ctx, set(params))
+  subst = dict(zip(params, args))
+  return (None,) + tuple(grads[p].substitute(subst) if p in grads else None for p in params)

 # ctx is grad_output
 pm_gradient = PatternMatcher([
@@ -1,6 +1,5 @@
 from __future__ import annotations
 import os, functools, platform, time, re, contextlib, operator, hashlib, pickle, sqlite3, tempfile, pathlib, string, ctypes, sys, gzip, getpass, gc
-from collections import defaultdict
 import subprocess, shutil, math, types, copyreg, inspect, importlib, decimal, itertools
 from dataclasses import dataclass, field
 from typing import ClassVar, Iterable, Any, TypeVar, Callable, Sequence, TypeGuard, Iterator, Generic, Generator, cast, overload
@@ -173,8 +172,7 @@ class ContextVar(Generic[T]):
    assert isinstance(self.value, str)
    return [getattr(obj, x) if obj else x for x in self.value.split(',') if x]

-DEBUG, BEAM, NOOPT = ContextVar("DEBUG", 0), ContextVar("BEAM", 0), ContextVar("NOOPT", 0)
-IMAGE, FLOAT16 = ContextVar("IMAGE", 0), ContextVar("FLOAT16", 0)
+DEBUG, IMAGE, BEAM, NOOPT = ContextVar("DEBUG", 0), ContextVar("IMAGE", 0), ContextVar("BEAM", 0), ContextVar("NOOPT", 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)
 USE_TC, TC_SELECT, TC_OPT, AMX = ContextVar("TC", 1), ContextVar("TC_SELECT", -1), ContextVar("TC_OPT", 0), ContextVar("AMX", 0)
@@ -233,7 +231,6 @@ class GlobalCounters:
  time_sum_s: ClassVar[float] = 0.0
  kernel_count: ClassVar[int] = 0
  mem_used: ClassVar[int] = 0   # NOTE: this is not reset
-  mem_used_per_device: ClassVar[defaultdict] = defaultdict(int)   # NOTE: this is not reset
  @staticmethod
  def reset(): GlobalCounters.global_ops, GlobalCounters.global_mem, GlobalCounters.time_sum_s, GlobalCounters.kernel_count = 0,0,0.0,0

@@ -2,6 +2,7 @@ from __future__ import annotations
 import math
 from tinygrad.tensor import Tensor
 from tinygrad.dtype import dtypes
+from tinygrad.device import is_dtype_supported
 from tinygrad.helpers import prod, make_tuple, flatten, USE_ATOMICS
 from tinygrad.nn import optim, state, datasets  # noqa: F401

@@ -35,7 +36,7 @@ class BatchNorm:
    self.weight: Tensor|None = Tensor.ones(sz) if affine else None
    self.bias: Tensor|None = Tensor.zeros(sz) if affine else None

-    self.num_batches_tracked = Tensor.zeros(dtype='long', requires_grad=False)
+    self.num_batches_tracked = Tensor.zeros(dtype='long' if is_dtype_supported(dtypes.long) else 'int', requires_grad=False)
    if track_running_stats: self.running_mean, self.running_var = Tensor.zeros(sz, requires_grad=False), Tensor.ones(sz, requires_grad=False)

  def calc_stats(self, x:Tensor) -> tuple[Tensor, Tensor]:
@@ -17,7 +17,6 @@ class Optimizer:
    assert len(self.params) != 0, "optimizer must have at least one param"
    self.buffers: list[Tensor] = dedup([x for x in params if not x.requires_grad])   # buffers are still realized
    self.device = device or self.params[0].device
-    self.param_dtype = to_dtype(getenv("OPTIM_DTYPE", "float32"))
    self.fused = fused
    # store lr in at least float32 precision
    self.lr = Tensor(lr if getenv("CONST_LR") else [lr], requires_grad=False, device=self.device,
@@ -25,9 +24,10 @@ class Optimizer:
    if self.fused: self.pos_params = list(itertools.accumulate(self.params, lambda x,y: x+y.numel(), initial=0))

  def _new_optim_param(self) -> list[Tensor]:
-    if self.fused: return [Tensor.zeros(self.pos_params[-1], dtype=self.param_dtype, device=self.device, requires_grad=False)]
-    if isinstance(self.device, tuple): return [Tensor.zeros_like(t, dtype=self.param_dtype, requires_grad=False) for t in self.params]
-    else: return [Tensor.zeros(t.shape, dtype=self.param_dtype, device=self.device, requires_grad=False) for t in self.params]
+    param_dtype = to_dtype(getenv("OPTIM_DTYPE", "float32"))
+    if self.fused: return [Tensor.zeros(self.pos_params[-1], dtype=param_dtype, device=self.device, requires_grad=False)]
+    if isinstance(self.device, tuple): return [Tensor.zeros_like(t, dtype=param_dtype, requires_grad=False) for t in self.params]
+    else: return [Tensor.zeros(t.shape, dtype=param_dtype, device=self.device, requires_grad=False) for t in self.params]

  def zero_grad(self):
    """
@@ -78,7 +78,7 @@ def safe_save(tensors:dict[str, Tensor], fn:str, metadata:dict[str, Any]|None=No
  j += "\x20"*(round_up(len(j),8)-len(j))
  pathlib.Path(fn).unlink(missing_ok=True)
  t = Tensor.empty(8+len(j)+offset, dtype=dtypes.uint8, device=f"disk:{fn}")
-  t[0:8].bitcast(dtypes.int64).assign([len(j)])
+  t[0:8].assign(Tensor([len(j)], dtype=dtypes.int64, device="CPU").bitcast(dtypes.uint8))
  t[8:8+len(j)].assign(list(j.encode('utf-8')))
  for k,v in safe_load(t).items(): v.assign(tensors[k])

@@ -304,7 +304,7 @@ def ggml_data_to_tensor(t: Tensor, n: int, ggml_type: int) -> Tensor:

  # native types
  if (dtype := { 0: dtypes.float32, 1: dtypes.float16, 16: dtypes.int8, 17: dtypes.int16, 18: dtypes.int32 }.get(ggml_type)) is not None:
-    return t[:dtype.itemsize * n].contiguous().bitcast(dtype)
+    return t[:dtype.itemsize * n].bitcast(dtype)

  def q_to_uint8(t: Tensor, b: int) -> Tensor:
    # TODO: rewrite with arange?
@@ -313,7 +313,7 @@ def ggml_data_to_tensor(t: Tensor, n: int, ggml_type: int) -> Tensor:

  # map to (number of elements, number of bytes)
  if (nelements_nbytes := { 2: (32, 18), 3: (32, 20), 8: (32, 34), 12: (256, 144), 14: (256, 210), 39: (32, 17) }.get(ggml_type)) is not None:
-    blocks = t[:(n//nelements_nbytes[0])*nelements_nbytes[1]].reshape((-1, nelements_nbytes[1])).contiguous()
+    blocks = t[:(n//nelements_nbytes[0])*nelements_nbytes[1]].reshape((-1, nelements_nbytes[1]))
    if ggml_type == 2: return (q_to_uint8(blocks[:,2:], 4).bitcast(dtypes.int8) - 8) * blocks[:,:2].bitcast(dtypes.float16).cast(dtypes.float32)
    if ggml_type == 3:
      d, m = (blocks[:,s:s+2].bitcast(dtypes.float16).cast(dtypes.float32) for s in [ 0, 2 ])
@@ -46,7 +46,6 @@ class InstOp(Enum):
  SMEM = 0x1
  JUMP = 0x3              # branch taken
  JUMP_NO = 0x4           # branch not taken
-  CALL = 0x5              # s_call_b64
  MESSAGE = 0x9
  VALU_TRANS = 0xb        # transcendental: exp, log, rcp, sqrt, sin, cos
  VALU_64_SHIFT = 0xd     # 64-bit shifts: lshl, lshr, ashr
@@ -73,10 +72,8 @@ class InstOp(Enum):

  # LDS ops on traced SIMD
  LDS_LOAD = 0x29
-  LDS_ATOMIC = 0x2a        # ds_append, ds_consume, ds_store_addtid_b32
  LDS_STORE = 0x2b
  LDS_STORE_64 = 0x2c
-  LDS_STORE_96 = 0x2d
  LDS_STORE_128 = 0x2e

  # Memory ops on other SIMD (0x5x range)
@@ -102,36 +99,19 @@ class InstOp(Enum):

 class InstOpRDNA4(Enum):
  """SQTT instruction operation types for RDNA4 (gfx1200). Different encoding from RDNA3."""
+  # TODO: we need to do discovery of all of these from instructions
  SALU = 0x0
-  SMEM = 0x1
-  JUMP = 0x3
-  JUMP_NO = 0x4
-  JUMP_UNCOND = 0x5
-  MESSAGE = 0x9
-  VALU_TRANS = 0xb
-  VALU_B2 = 0xd
-  VALU_B4 = 0xe
-  VINTERP = 0x12
-  VMEM_RD_1 = 0x21
-  VMEM_WR_2 = 0x24
-  VMEM_WR_3 = 0x25
-  VMEM_WR_4 = 0x26
-  VMEM_WR_5 = 0x27
-  VMEM_WR_6 = 0x28
-  LDS_RD = 0x29
-  LDS_WR_1 = 0x2a
-  LDS_WR_2 = 0x2b
-  LDS_WR_3 = 0x2c
-  LDS_WR_4 = 0x2d
-  LDS_WR_5 = 0x2e
-  WMMA_8 = 0x8c
-  WMMA_16 = 0x8d
-  VALU_DPFP = 0x92
-  SALU_FLOAT3 = 0x98
-  VALU_SCL_TRANS = 0x99
-  SALU_2 = 0x9b
-  SALU_5 = 0x9c
-  OTHER_VMEM = 0xc1
+  JUMP = 0x1
+  NEXT = 0x2
+  MESSAGE = 0x4
+  VALU_64 = 0x6
+  VALU_WMMA = 0x46
+  VMEM = 0x10
+  VMEM_128 = 0x11
+  VMEM_STORE = 0x12
+  VMEM_STORE_128 = 0x14
+  OTHER_VMEM = 0x5e
+  OTHER_VMEM_STORE = 0x60

 # ═══════════════════════════════════════════════════════════════════════════════
 # PACKET TYPE BASE CLASS
@@ -167,6 +147,11 @@ class TS_DELTA_S8_W3(PacketType):
  delta = bits[10:8]
  _padding = bits[63:11]

+class TS_DELTA_S8_W3_RDNA4(PacketType):  # Layout 4: 64->72 bits
+  encoding = bits[6:0] == 0b0100001
+  delta = bits[10:8]
+  _padding = bits[71:11]
+
 class TS_DELTA_S5_W3(PacketType):
  encoding = bits[4:0] == 0b00110
  delta = bits[7:5]
@@ -356,9 +341,14 @@ class INST(PacketType):
 class INST_RDNA4(PacketType):  # Layout 4: different delta position and InstOp encoding
  encoding = bits[2:0] == 0b010
  delta = bits[5:3]
-  w64h = bits[6:6]
-  wave = bits[11:7]
-  op = bits[19:12].enum(InstOpRDNA4)
+  flag1 = bits[6:6]
+  flag2 = bits[7:7]
+  wave_pair = bits[11:8]
+  flag3 = bits[12:12]
+  op = bits[19:13].enum(InstOpRDNA4)
+  # INST_RDNA4 wave_pair field (4 bits) addresses wave pairs, flag2 selects even/odd wave
+  @property
+  def wave(self): return self.wave_pair * 2 + self.flag2

 class UTILCTR(PacketType):
  encoding = bits[6:0] == 0b0110001
@@ -373,7 +363,7 @@ PACKET_TYPES_RDNA3: dict[int, type[PacketType]] = {
 }
 PACKET_TYPES_RDNA4: dict[int, type[PacketType]] = {
  **PACKET_TYPES_RDNA3,
-  9: WAVESTART_RDNA4, 10: TS_DELTA_S5_W2_RDNA4, 11: WAVEALLOC_RDNA4,
+  7: TS_DELTA_S8_W3_RDNA4, 9: WAVESTART_RDNA4, 10: TS_DELTA_S5_W2_RDNA4, 11: WAVEALLOC_RDNA4,
  12: TS_DELTA_S5_W3_RDNA4, 13: PERF_RDNA4, 22: TS_DELTA_OR_MARK_RDNA4, 24: INST_RDNA4,
 }

@@ -546,9 +536,8 @@ def decode(data: bytes) -> Iterator[PacketType]:
    if nib_off: reg, pos = (reg >> 4) | ((data[pos] >> 4) << 60), pos + 1
    # 2. read all full bytes at once
    if (byte_count := need >> 1):
-      read_bytes = min(byte_count, 8)
-      chunk = int.from_bytes(data[pos:pos + read_bytes], 'little')
-      reg, pos = (reg >> (read_bytes * 8)) | (chunk << (64 - read_bytes * 8)), pos + byte_count
+      chunk = int.from_bytes(data[pos:pos + byte_count], 'little')
+      reg, pos = (reg >> (byte_count * 8)) | (chunk << (64 - byte_count * 8)), pos + byte_count
    # 3. if odd, read low nibble
    if (nib_off := need & 1): reg = (reg >> 4) | ((data[pos] & 0xF) << 60)

@@ -630,9 +619,9 @@ def map_insts(data:bytes, lib:bytes, target:str) -> Iterator[tuple[PacketType, I
      # identify a branch instruction, only used for asserts
      branch_inst = inst if "BRANCH" in inst_op else None
      if branch_inst is not None:
-        assert isinstance(p, (INST, INST_RDNA4)) and p.op.name in {"JUMP_NO", "JUMP", "JUMP_UNCOND"}, f"branch can only be folowed by JUMP, got {p}"
+        assert isinstance(p, (INST, INST_RDNA4)) and p.op.name in {"JUMP_NO", "JUMP", "NEXT"}, f"branch can only be folowed by JUMP, got {p}"
      # JUMP handling
-      if (isinstance(p, INST) and p.op is InstOp.JUMP) or (isinstance(p, INST_RDNA4) and p.op is InstOpRDNA4.JUMP):
+      if (isinstance(p, INST) and p.op is InstOp.JUMP) or (isinstance(p, INST_RDNA4) and branch_inst is not None and p.flag3):
        simm16 = getattr(branch_inst, 'simm16')
        assert branch_inst is not None and simm16 is not None, f"JUMP packet must map to a branch instruction, got {inst}"
        x = simm16 & 0xffff
@@ -661,7 +650,7 @@ def format_packet(p) -> str:
  name = type(p).__name__
  if isinstance(p, (INST, INST_RDNA4)):
    op_name = p.op.name if isinstance(p.op, (InstOp, InstOpRDNA4)) else f"0x{p.op:02x}"
-    fields = f"wave={p.wave} op={op_name}" + ((" flag1" if p.flag1 else "") + (" flag2" if p.flag2 else "") if isinstance(p, INST) else "")
+    fields = f"wave={p.wave} op={op_name}" + (" flag1" if p.flag1 else "") + (" flag2" if p.flag2 else "")
  elif isinstance(p, VALUINST): fields = f"wave={p.wave}" + (" flag" if p.flag else "")
  elif isinstance(p, ALUEXEC): fields = f"src={p.src.name if isinstance(p.src, AluSrc) else p.src}"
  elif isinstance(p, VMEMEXEC): fields = f"src={p.src.name if isinstance(p.src, MemSrc) else p.src}"
@@ -677,19 +666,18 @@ def print_packets(packets) -> None:
  from tinygrad.helpers import getenv
  skip = {"NOP", "TS_DELTA_SHORT", "TS_WAVE_STATE", "TS_DELTA_OR_MARK",
          "TS_DELTA_S5_W2", "TS_DELTA_S5_W3", "TS_DELTA_S8_W3", "REG", "EVENT"} if not getenv("NOSKIP") else {"NOP"}
-  for data in packets:
-    p, inst = data if isinstance(data, tuple) else (data, None)
-    if type(p).__name__.replace("_RDNA4", "") not in skip: print(format_packet(p), f"inst={inst.inst}" if inst is not None else '')
+  for p in packets:
+    if type(p).__name__.replace("_RDNA4", "") not in skip: print(format_packet(p))

 if __name__ == "__main__":
  import sys, pickle
-  from tinygrad.helpers import temp
-  with open(temp("profile.pkl", append_user=True) if len(sys.argv) < 2 else sys.argv[1], "rb") as f:
+  if len(sys.argv) < 2:
+    print("Usage: python sqtt.py <pkl_file>")
+    sys.exit(1)
+  with open(sys.argv[1], "rb") as f:
    data = pickle.load(f)
-  prg_events = {e.tag: e for e in data if type(e).__name__ == "ProfileProgramEvent" and e.tag is not None}
+  prg_names = {e.tag: e.name for e in data if type(e).__name__ == "ProfileProgramEvent" and e.tag is not None}
  sqtt_events = [e for e in data if type(e).__name__ == "ProfileSQTTEvent"]
-  dev_targets = {e.device:f"gfx{e.props['gfx_target_version']//1000}" for e in data if type(e).__name__ == "ProfileDeviceEvent" and e.props}
  for i, event in enumerate(sqtt_events):
-    prg = prg_events.get(event.kern)
-    print(f"\n=== event {i} {prg.name if prg is not None else ''} ===")
-    print_packets(map_insts(event.blob, prg.lib, dev_targets[prg.device]) if prg is not None else decode(event.blob))
+    print(f"\n=== event {i} {prg_names.get(event.kern, '')} ===")
+    print_packets(decode(event.blob))
@@ -598,10 +598,9 @@ class AMDProgram(HCQProgram):
                     base=self.lib_gpu.va_addr)
    weakref.finalize(self, self._fini, self.dev, self.lib_gpu, buf_spec)

-  def __call__(self, *bufs, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1), vals:tuple[int|None, ...]=(),
-               wait=False, timeout:int|None=None):
+  def __call__(self, *bufs, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1), vals:tuple[int|None, ...]=(), wait=False):
    if self.dev.sqtt_enabled: cast(AMDComputeQueue, self.dev.hw_compute_queue_t()).sqtt_start(self.dev.sqtt_buffers).submit(self.dev)
-    res = super().__call__(*bufs, global_size=global_size, local_size=local_size, vals=vals, wait=wait, timeout=timeout)
+    res = super().__call__(*bufs, global_size=global_size, local_size=local_size, vals=vals, wait=wait)
    if self.dev.pmc_enabled:
      cast(AMDComputeQueue, self.dev.hw_compute_queue_t()).pmc_read(self.dev.pmc_buffer, self.dev.pmc_sched) \
                                                          .signal(self.dev.timeline_signal, self.dev.next_timeline()).submit(self.dev)
@@ -858,33 +857,31 @@ class PCIIface(PCIIfaceBase):

    rcvr_params: tuple
    if queue_type == kfd.KFD_IOC_QUEUE_TYPE_SDMA:
-      doorbell_index = self.dev_impl.sdma.setup_ring(*(rcvr_params:=(ring.va_addr, ring.size, gart.va_addr+rptr, gart.va_addr+wptr, idx)))
+      pv, doorbell_index = self.dev_impl.sdma.setup_ring(*(rcvr_params:=(ring.va_addr, ring.size, gart.va_addr+rptr, gart.va_addr+wptr, idx)))
    else:
-      doorbell_index = self.dev_impl.gfx.setup_ring(*(rcvr_params:=(ring.va_addr, ring.size, gart.va_addr+rptr, gart.va_addr+wptr,
+      pv, doorbell_index = self.dev_impl.gfx.setup_ring(*(rcvr_params:=(ring.va_addr, ring.size, gart.va_addr+rptr, gart.va_addr+wptr,
        eop_buffer.va_addr, eop_buffer.size, is_aql:=(queue_type==kfd.KFD_IOC_QUEUE_TYPE_COMPUTE_AQL), is_aql)))

-    return AMDQueueDesc(ring=ring.cpu_view().view(fmt='I'), doorbell=self.dev_impl.doorbell64.view(doorbell_index * 8, 8, fmt='Q'), put_value=0,
+    return AMDQueueDesc(ring=ring.cpu_view().view(fmt='I'), doorbell=self.dev_impl.doorbell64.view(doorbell_index * 8, 8, fmt='Q'), put_value=pv,
      read_ptr=gart.cpu_view().view(offset=rptr, size=8, fmt='Q'), write_ptr=gart.cpu_view().view(offset=wptr, size=8, fmt='Q'), params=rcvr_params)

-  def _collect_faults(self, reset=False):
-    devs:list[AMDDevice] = [d for pg in HCQCompiled.peer_groups.values() for d in pg if isinstance(d, AMDDevice) and d.is_am()]
-    for d in devs:
-      d.iface.dev_impl.ih.interrupt_handler()
-      if reset and d.iface.dev_impl.recover(force=d.error_state is not None):
-        d.compute_queue.put_value = d.compute_queue.read_ptr[0] = d.compute_queue.write_ptr[0] = 0
-        d.iface.dev_impl.gfx.setup_ring(*d.compute_queue.params)
-        d.timeline_signal.value = d.timeline_value - 1
-        d.error_state = None
-
  def sleep(self, timeout):
    if hasattr(self.pci_dev, 'irq_poller') and self.pci_dev.irq_poller is not None and (events_cnt:=len(self.pci_dev.irq_poller.poll(timeout))):
      self.pci_dev.irq_fd.read(8 * events_cnt)
-    self._collect_faults()
+    self.dev_impl.ih.interrupt_handler()
    if self.dev_impl.is_err_state: raise RuntimeError("Device is in error state")

  def on_device_hang(self):
-    self._collect_faults(reset=True)
-    raise RuntimeError("Device hang detected")
+    devs:list[AMDDevice] = [d for pg in HCQCompiled.peer_groups.values() for d in pg if isinstance(d, AMDDevice) and d.is_am()]
+    for d in devs: d.iface.dev_impl.ih.interrupt_handler()
+    faults = [f for d in devs if (f:=d.iface.dev_impl.gmc.check_fault())]
+    for d in devs:
+      if d.iface.dev_impl.recover():
+        d.compute_queue.put_value, _ = d.iface.dev_impl.gfx.setup_ring(*d.compute_queue.params)
+        d.compute_queue.read_ptr[0] = d.compute_queue.write_ptr[0] = d.compute_queue.put_value
+        d.timeline_signal.value = d.timeline_value - 1
+        d.error_state = None
+    raise RuntimeError(f"Device hang detected: {'; '.join(faults)}" if faults else "Device hang detected")

  def device_fini(self): self.dev_impl.fini()

@@ -978,8 +975,7 @@ class AMDDevice(HCQCompiled):
    super().__init__(device, AMDAllocator(self), compilers, functools.partial(AMDProgram, self), AMDSignal,
                     functools.partial(AMDComputeAQLQueue if self.is_aql else AMDComputeQueue, self),
                     functools.partial(AMDCopyQueue, self, max_copy_size=self.max_copy_size) if self.has_sdma_queue else None,
-                     kernargs_size=(8 << 10) if self.is_usb() else (16 << 20), sigalloc_size=0x100 if self.is_usb() else 0x1000,
-                     can_recover=self.is_am())
+                     kernargs_size=(8 << 10) if self.is_usb() else (16 << 20), sigalloc_size=0x100 if self.is_usb() else 0x1000)

    # Scratch setup
    self.max_private_segment_size = 0
@@ -54,7 +54,7 @@ class CLProgram:
    except (TypeError, AttributeError): pass

  def __call__(self, *bufs:tuple[cl.cl_mem, BufferSpec], global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]|None=None,
-               vals:tuple[int, ...]=(), wait=False, **kw) -> float|None:
+               vals:tuple[int, ...]=(), wait=False) -> float|None:
    i = 0
    for i,(b,_) in enumerate(bufs):
      for real_i, dt in self.arg_dtypes[i]:
@@ -51,7 +51,7 @@ class CUDAProgram:
  @suppress_finalizing
  def __del__(self): check(cuda.cuModuleUnload(self.module))

-  def __call__(self, *args, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1), vals:tuple[int, ...]=(), wait=False, **kw):
+  def __call__(self, *args, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1), vals:tuple[int, ...]=(), wait=False):
    check(cuda.cuCtxSetCurrent(self.dev.context))
    if not hasattr(self, "vargs"):
      self.c_args, self.vargs = encode_args(args, vals)
@@ -84,7 +84,7 @@ class DSPProgram:
  def __init__(self, dev:DSPDevice, name:str, lib:bytes, **kwargs):
    self.dev, self.lib = dev, lib

-  def __call__(self, *bufs, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1), vals:tuple[int, ...]=(), wait=False, **kw):
+  def __call__(self, *bufs, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1), vals:tuple[int, ...]=(), wait=False):
    if len(bufs) >= 16: raise RuntimeError(f"Too many buffers to execute: {len(bufs)}")

    pra, fds, attrs, _ = rpc_prep_args(ins=[var_vals_mv:=memoryview(bytearray((len(bufs)+len(vals))*4)), off_mv:=memoryview(bytearray(len(bufs)*4))],
@@ -293,7 +293,7 @@ class MockDSPRenderer(DSPRenderer):

 class MockDSPProgram:
  def __init__(self, name:str, lib:bytes, **kwargs): self.lib = lib
-  def __call__(self, *bufs, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1), vals:tuple[int, ...]=(), wait=False, **kw):
+  def __call__(self, *bufs, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1), vals:tuple[int, ...]=(), wait=False):
    with tempfile.NamedTemporaryFile(suffix=".out") as dsp_lib:
      dsp_lib.write(self.lib)
      dsp_lib.flush()
@@ -32,7 +32,7 @@ class HIPProgram:
  def __del__(self):
    if hasattr(self, 'module'): check(hip.hipModuleUnload(self.module))

-  def __call__(self, *args, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1), vals:tuple[int, ...]=(), wait=False, **kw):
+  def __call__(self, *args, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1), vals:tuple[int, ...]=(), wait=False):
    check(hip.hipSetDevice(self.dev.device_id))
    if not hasattr(self, "vargs"):
      fields = [(f'f{i}', hip.hipDeviceptr_t, i*8) for i in range(len(args))] + [(f'v{i}', ctypes.c_int, len(args)*8+i*4) for i in range(len(vals))]
@@ -123,7 +123,7 @@ class MetalProgram:
    # cache these msg calls
    self.max_total_threads: int = self.pipeline_state.maxTotalThreadsPerThreadgroup()

-  def __call__(self, *bufs, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1), vals:tuple[int, ...]=(), wait=False, **kw):
+  def __call__(self, *bufs, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1), vals:tuple[int, ...]=(), wait=False):
    if prod(local_size) > self.max_total_threads:
      exec_width = self.pipeline_state.threadExecutionWidth()
      memory_length = self.pipeline_state.staticThreadgroupMemoryLength()
@@ -15,7 +15,7 @@ class NullRenderer(CStyleLanguage):

 class NullProgram:
  def __init__(self, device:str, name:str, lib:bytes, *args, **kwargs): self.device, self.name = device, name
-  def __call__(self, *bufs, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1), vals:tuple[int, ...]=(), wait=False, **kw):
+  def __call__(self, *bufs, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1), vals:tuple[int, ...]=(), wait=False):
    with cpu_profile(self.name, self.device): return 1e-3

 class NullAllocator(Allocator['NullDevice']):
@@ -312,13 +312,12 @@ class NVProgram(HCQProgram):
      yield typ, param, sh.content[start_off+4:start_off+sz+4] if typ == 0x4 else sz
      start_off += (sz if typ == 0x4 else 0) + 4

-  def __call__(self, *bufs, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1), vals:tuple[int|None, ...]=(),
-               wait=False, timeout:int|None=None):
+  def __call__(self, *bufs, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1), vals:tuple[int|None, ...]=(), wait=False):
    if prod(local_size) > 1024 or self.max_threads < prod(local_size) or self.lcmem_usage > cast(NVDevice, self.dev).slm_per_thread:
      raise RuntimeError(f"Too many resources requested for launch, {prod(local_size)=}, {self.max_threads=}")
    if any(cur > mx for cur,mx in zip(global_size, [2147483647, 65535, 65535])) or any(cur > mx for cur,mx in zip(local_size, [1024, 1024, 64])):
      raise RuntimeError(f"Invalid global/local dims {global_size=}, {local_size=}")
-    res = super().__call__(*bufs, global_size=global_size, local_size=local_size, vals=vals, wait=wait, timeout=timeout)
+    res = super().__call__(*bufs, global_size=global_size, local_size=local_size, vals=vals, wait=wait)
    if self.dev.pma_enabled:
      self.dev.synchronize()
      if pma_blob:=self.dev._prof_readback():
@@ -41,7 +41,7 @@ def generic_wmma_helper(inp, warp_size, WARP_THREADS, K, NUM_A, NUM_B, NUM_C, a_
 class PythonProgram:
  def __init__(self, name:str, lib:bytes, **kwargs):
    self.uops: list[tuple[Ops, DType, list[int], Any]] = pickle.loads(lib)
-  def __call__(self, *bufs, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1), vals:tuple[int, ...]=(), wait=False, **kw):
+  def __call__(self, *bufs, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1), vals:tuple[int, ...]=(), wait=False):
    st = time.perf_counter()
    warp = list(itertools.product(*[range(x) for x in local_size[::-1]]))
    warp_size = len(warp)
@@ -266,8 +266,7 @@ class QCOMProgram(HCQProgram):
    super().__init__(QCOMArgsState, self.dev, self.name, kernargs_alloc_size=kernargs_alloc_size)
    weakref.finalize(self, self._fini, self.dev, self.lib_gpu, buf_spec)

-  def __call__(self, *bufs, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1),
-               vals:tuple[int|None, ...]=(), wait=False, **kw):
+  def __call__(self, *bufs, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1), vals:tuple[int|None, ...]=(), wait=False):
    if self.max_threads < prod(local_size): raise RuntimeError("Too many resources requested for launch")
    if any(g*l>mx for g,l,mx in zip(global_size, local_size, [65536, 65536, 65536])) and any(l>mx for l,mx in zip(local_size, [1024, 1024, 1024])):
      raise RuntimeError(f"Invalid global/local dims {global_size=}, {local_size=}")
@@ -90,7 +90,7 @@ class WebGPUProgram:

    self.name, self.lib, self.prg = name, lib, shader_module
  def __call__(self, *bufs:WGPUBufPtr, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1),
-               vals:tuple[int, ...]=(), wait=False, **kw) -> float|None:
+               vals:tuple[int, ...]=(), wait=False) -> float|None:
    wait = wait and self.timestamp_supported
    tmp_bufs = [*bufs]
    buf_patch = False
@@ -193,7 +193,7 @@ class AMDev(PCIDevImplBase):
    if DEBUG >= 2: print(f"am {self.devfmt}: boot done")

  def init_sw(self, smi_dev=False):
-    self.smi_dev, self.is_err_state = smi_dev, False
+    self.smi_dev, self.is_err_state, self.has_aql_queue = smi_dev, False, False

    # Memory manager & firmware
    self.mm = AMMemoryManager(self, self.vram_size - self.reserved_vram_size, boot_size=(32 << 20), pt_t=AMPageTableEntry, va_shifts=[12, 21, 30, 39],
@@ -225,13 +225,13 @@ class AMDev(PCIDevImplBase):
    self.ih.interrupt_handler()
    self.reg("regSCRATCH_REG6").write(self.is_err_state) # set finalized state.

-  def recover(self, force=False) -> bool:
-    if not force and not self.is_err_state: return False
-    if DEBUG >= 3: print(f"am {self.devfmt}: Start recovery")
+  def recover(self) -> bool:
+    if (self.has_aql_queue and self.is_hive()) or not self.is_err_state: return False # TODO: support aql queue recovery on hive
+    if DEBUG >= 2: print(f"am {self.devfmt}: Start recovery")
    self.ih.interrupt_handler()
    self.gfx.reset_mec()
    self.is_err_state = False
-    if DEBUG >= 3: print(f"am {self.devfmt}: Recovery complete")
+    if DEBUG >= 2: print(f"am {self.devfmt}: Recovery complete")
    return True

  def is_hive(self) -> bool: return self.gmc.xgmi_seg_sz > 0
@@ -243,19 +243,19 @@ class AMDev(PCIDevImplBase):
  def reg(self, reg:str) -> AMRegister: return self.__dict__[reg]

  def rreg(self, reg:int) -> int:
-    val = self.indirect_rreg(reg) if reg >= len(self.mmio) else self.mmio[reg]
+    val = self.indirect_rreg(reg) if reg > len(self.mmio) else self.mmio[reg]
    if AM_DEBUG >= 4 and getattr(self, '_prev_rreg', None) != (reg, val): print(f"am {self.devfmt}: Reading register {reg:#x} with value {val:#x}")
    self._prev_rreg = (reg, val)
    return val

  def wreg(self, reg:int, val:int):
    if AM_DEBUG >= 4: print(f"am {self.devfmt}: Writing register {reg:#x} with value {val:#x}")
-    if reg >= len(self.mmio): self.indirect_wreg(reg, val)
+    if reg > len(self.mmio): self.indirect_wreg(reg, val)
    else: self.mmio[reg] = val

  def wreg_pair(self, reg_base:str, lo_suffix:str, hi_suffix:str, val:int, inst:int=0):
-    self.reg(f"{reg_base}{lo_suffix}").write(lo32(val), inst=inst)
-    self.reg(f"{reg_base}{hi_suffix}").write(hi32(val), inst=inst)
+    self.reg(f"{reg_base}{lo_suffix}").write(val & 0xffffffff, inst=inst)
+    self.reg(f"{reg_base}{hi_suffix}").write(val >> 32, inst=inst)

  def indirect_rreg(self, reg:int) -> int:
    self.reg("regBIF_BX_PF0_RSMU_INDEX").write(reg * 4)
@@ -268,9 +268,9 @@ class AMDev(PCIDevImplBase):
  def indirect_wreg_pcie(self, reg:int, val:int, aid:int=0):
    reg_addr = reg * 4 + ((((aid & 0b11) << 32) | (1 << 34)) if aid > 0 else 0)
    self.reg("regBIF_BX0_PCIE_INDEX2").write(lo32(reg_addr))
-    if hi32(reg_addr) > 0: self.reg("regBIF_BX0_PCIE_INDEX2_HI").write(hi32(reg_addr) & 0xff)
+    if reg_addr >> 32: self.reg("regBIF_BX0_PCIE_INDEX2_HI").write(hi32(reg_addr) & 0xff)
    self.reg("regBIF_BX0_PCIE_DATA2").write(val)
-    if hi32(reg_addr) > 0: self.reg("regBIF_BX0_PCIE_INDEX2_HI").write(0)
+    if reg_addr >> 32: self.reg("regBIF_BX0_PCIE_INDEX2_HI").write(0)

  def _read_vram(self, addr, size) -> bytes:
    assert addr % 4 == 0 and size % 4 == 0, f"Invalid address {addr:#x} or size {size:#x}"
@@ -25,7 +25,7 @@ class AM_SOC(AM_IP):
      return {getattr(am, k): k[off+9:] for k in dir(am) if k.startswith(f'{pref}_{self.adev.ip_ver[hwip][0]}') and (off:=k.find('__SRCID__')) != -1}

    gfx_srcs, sdma_srcs = _ih_srcs('GFX', am.GC_HWIP), _ih_srcs('SDMA0', am.SDMA0_HWIP)
-    self.ih_srcs_names:dict[int, dict[int, str]] = {**{k: gfx_srcs for k in self.gfx_ih_clients}, **{k: sdma_srcs for k in self.sdma_ih_clients}}
+    self.ih_scrs_names:dict[int, dict[int, str]] = {**{k: gfx_srcs for k in self.gfx_ih_clients}, **{k: sdma_srcs for k in self.sdma_ih_clients}}

  def init_hw(self):
    if self.adev.ip_ver[am.NBIO_HWIP] in {(7,9,0), (7,9,1)}:
@@ -171,6 +171,12 @@ class AM_GMC(AM_IP):
    if self.adev.ip_ver[am.GC_HWIP] < (10,0,0): return (pte & am.AMDGPU_PDE_PTE) if pte_lv != am.AMDGPU_VM_PDB0 else not (pte & am.AMDGPU_PTE_TF)
    return pte & (am.AMDGPU_PDE_PTE_GFX12 if self.adev.ip_ver[am.GC_HWIP] >= (12,0,0) else am.AMDGPU_PDE_PTE)

+  def check_fault(self) -> str|None:
+    va = (self.adev.reg('regGCVM_L2_PROTECTION_FAULT_ADDR_HI32').read()<<32) | self.adev.reg('regGCVM_L2_PROTECTION_FAULT_ADDR_LO32').read()
+    if self.adev.reg(self.pf_status_reg("GC")).read():
+      return f"am {self.adev.devfmt}: GCVM_L2_PROTECTION_FAULT_STATUS: {self.adev.reg(self.pf_status_reg('GC')).read_bitfields()} {va<<12:#x}"
+    return None
+
 class AM_SMU(AM_IP):
  def init_sw(self):
    self.smu_mod = self.adev._ip_module("smu", am.MP1_HWIP, prever_prefix='v')
@@ -240,11 +246,10 @@ class AM_GFX(AM_IP):

  def init_hw(self):
    # Wait for RLC autoload to complete
-    wait_cond(lambda: self.adev.regCP_STAT.read() == 0 or self.adev.regRLC_RLCS_BOOTLOAD_STATUS.read_bitfields()['bootload_complete'] == 0,
-              value=True, msg="RLC autoload timeout")
+    while self.adev.regCP_STAT.read() != 0 and self.adev.regRLC_RLCS_BOOTLOAD_STATUS.read_bitfields()['bootload_complete'] != 0: pass

    self.adev.gmc.init_hub("GC", inst_cnt=self.xccs)
-    if self.adev.partial_boot: return self.reset_mec()
+    if self.adev.partial_boot: return

    self._config_mec()

@@ -286,26 +291,24 @@ class AM_GFX(AM_IP):
    self._enable_mec()

    # Set 1 partition
-    if self.xccs > 1: self.adev.psp._spatial_partition_cmd(1)
+    if self.xccs > 1 and not self.adev.partial_boot: self.adev.psp._spatial_partition_cmd(1)

  def fini_hw(self): self._dequeue_hqds()

  def reset_mec(self):
-    self._dequeue_hqds()
-
-    for xcc in range(self.xccs): self.adev.regGRBM_SOFT_RESET.write(soft_reset_cp=1, soft_reset_cpc=1, inst=xcc)
-    time.sleep(0.05)
-    for xcc in range(self.xccs): self.adev.regGRBM_SOFT_RESET.write(0x0, inst=xcc)
-
+    self._dequeue_hqds(reset=True)
    self._config_mec()
    self._enable_mec()

-  def setup_ring(self, ring_addr:int, ring_size:int, rptr_addr:int, wptr_addr:int, eop_addr:int, eop_size:int, idx:int, aql:bool) -> int:
+  def setup_ring(self, ring_addr:int, ring_size:int, rptr_addr:int, wptr_addr:int, eop_addr:int, eop_size:int, idx:int, aql:bool) -> tuple[int, int]:
+    self.adev.has_aql_queue |= aql
    pipe, queue, doorbell = idx // 4, idx % 4, am.AMDGPU_NAVI10_DOORBELL_MEC_RING0
+    self._grbm_select(me=1, pipe=pipe, queue=queue, inst=0)
+    restore_queue = aql and self.xccs > 1 and self.adev.partial_boot and (self.adev.regCP_HQD_ACTIVE.read(inst=0) & 1)
+    restore_ptr = (self.adev.regCP_HQD_PQ_WPTR_LO.read(inst=0) | (self.adev.regCP_HQD_PQ_WPTR_HI.read(inst=0) << 32)) if restore_queue else 0
+    if DEBUG >= 2 and restore_queue: print(f"am {self.adev.devfmt}: GFX queue already active, continuing from saved state {restore_ptr=:#x}.")

    for xcc in range(self.xccs if aql else 1):
-      self._grbm_select(me=1, pipe=pipe, queue=queue, inst=xcc)
-
      struct_t = getattr(am, f"struct_v{self.adev.ip_ver[am.GC_HWIP][0]}{'_compute' if self.adev.ip_ver[am.GC_HWIP][0] >= 10 else ''}_mqd")
      mqd_struct = struct_t(header=0xC0310800, cp_mqd_base_addr_lo=lo32(self.mqd_mc[queue] + 0x1000*xcc),
        cp_mqd_base_addr_hi=hi32(self.mqd_mc[queue] + 0x1000*xcc), cp_hqd_pipe_priority=0x2, cp_hqd_queue_priority=0xf, cp_hqd_quantum=0x111,
@@ -323,16 +326,26 @@ class AM_GFX(AM_IP):
        **({'compute_tg_chunk_size':1, 'compute_current_logic_xcc_id':xcc, 'cp_mqd_stride_size':0x1000} if aql and self.xccs > 1 else {}))
      for se in range(8 if self.adev.ip_ver[am.GC_HWIP][0] >= 10 else 4): setattr(mqd_struct, f'compute_static_thread_mgmt_se{se}', 0xffffffff)

-      self.adev.vram.view(self.mqd_paddr[queue] + 0x1000*xcc, ctypes.sizeof(mqd_struct))[:] = memoryview(mqd_struct).cast('B')
+      # Copy mqd into memory
+      self._grbm_select(me=1, pipe=pipe, queue=queue, inst=xcc)

-      mqd_st_mv = to_mv(ctypes.addressof(mqd_struct), ctypes.sizeof(mqd_struct)).cast('I')
-      for i, reg in enumerate(range(self.adev.regCP_MQD_BASE_ADDR.addr[xcc], self.adev.regCP_HQD_PQ_WPTR_HI.addr[xcc] + 1)):
-        self.adev.wreg(reg, mqd_st_mv[0x80 + i])
-      self.adev.regCP_HQD_ACTIVE.write(0x1, inst=xcc)
+      if restore_queue:
+        for r in [self.adev.regCP_HQD_PQ_RPTR_REPORT_ADDR, self.adev.regCP_HQD_EOP_BASE_ADDR, self.adev.regCP_HQD_EOP_BASE_ADDR_HI,
+                  self.adev.regCP_HQD_PQ_RPTR_REPORT_ADDR_HI, self.adev.regCP_HQD_PQ_WPTR_POLL_ADDR, self.adev.regCP_HQD_PQ_WPTR_POLL_ADDR_HI]:
+          val = memoryview(bytes(mqd_struct)).cast('I')[0x80 + (off:=r.addr[xcc] - self.adev.regCP_MQD_BASE_ADDR.addr[xcc])]
+          self.adev.vram.view(self.mqd_paddr[queue] + 0x1000*xcc, ctypes.sizeof(mqd_struct), fmt='I')[0x80 + off] = val
+          r.write(val, inst=xcc)
+      else:
+        self.adev.vram.view(self.mqd_paddr[queue] + 0x1000*xcc, ctypes.sizeof(mqd_struct))[:] = memoryview(mqd_struct).cast('B')
+
+        mqd_st_mv = to_mv(ctypes.addressof(mqd_struct), ctypes.sizeof(mqd_struct)).cast('I')
+        for i, reg in enumerate(range(self.adev.regCP_MQD_BASE_ADDR.addr[xcc], self.adev.regCP_HQD_PQ_WPTR_HI.addr[xcc] + 1)):
+          self.adev.wreg(reg, mqd_st_mv[0x80 + i])
+        self.adev.regCP_HQD_ACTIVE.write(0x1, inst=xcc)

      self.adev.gmc.flush_hdp()
      self._grbm_select(inst=xcc)
-    return doorbell
+    return restore_ptr // 16, doorbell

  def set_clockgating_state(self):
    if hasattr(self.adev, 'regMM_ATC_L2_MISC_CG'): self.adev.regMM_ATC_L2_MISC_CG.write(enable=1, mem_ls_enable=1)
@@ -383,13 +396,15 @@ class AM_GFX(AM_IP):
      if self.adev.ip_ver[am.GC_HWIP] >= (10,0,0):
        _config_helper(eng_name="MEC", cntl_reg="MEC_RS64", eng_reg="MEC_RS64", pipe_cnt=1, me=1, xcc=xcc)

-  def _dequeue_hqds(self):
-    for q in range(2):
+  def _dequeue_hqds(self, reset=False):
+    # NOTE: For aqls with xccs (queue=1), will continue from the saved state.
+    for q in range(2 if self.xccs == 1 else 1):
      for xcc in range(self.xccs):
        self._grbm_select(me=1, pipe=0, queue=q, inst=xcc)
        if self.adev.regCP_HQD_ACTIVE.read(inst=xcc) & 1:
          self.adev.regCP_HQD_DEQUEUE_REQUEST.write(0x2, inst=xcc) # 1 - DRAIN_PIPE; 2 - RESET_WAVES
-          if not self.adev.is_err_state: wait_cond(lambda: self.adev.regCP_HQD_ACTIVE.read(inst=xcc) & 1, value=0, msg="HQD dequeue timeout")
+          if reset: self.adev.regSPI_COMPUTE_QUEUE_RESET.write(1, inst=xcc)
+          else: wait_cond(lambda: self.adev.regCP_HQD_ACTIVE.read(inst=xcc) & 1, value=0, msg="HQD dequeue timeout")
    self._grbm_select()

 class AM_IH(AM_IP):
@@ -436,7 +451,7 @@ class AM_IH(AM_IP):
        [getattr(am, f'SOC15_{n}_FROM_IH_ENTRY')(entry) for n in ['CLIENT_ID', 'SOURCE_ID', 'RING_ID', 'VMID', 'VMID_TYPE', 'PASID', 'NODEID']]
      ctx = [getattr(am, f'SOC15_CONTEXT_ID{i}_FROM_IH_ENTRY')(entry) for i in range(4)]

-      src_name = self.adev.soc.ih_srcs_names.get(client, {}).get(src, '')
+      src_name = self.adev.soc.ih_scrs_names.get(client, {}).get(src, '')
      print(f"am {self.adev.devfmt}: IH ({rptr:#x}/{wptr['offset']:#x}) client={self.adev.soc.ih_clients.get(client)} src={src_name}({src}) "
            f"ring={ring_id} vmid={vmid}({vmid_type}) pasid={pasid} node={node} ctx=[{ctx[0]:#x}, {ctx[1]:#x}, {ctx[2]:#x}, {ctx[3]:#x}]")

@@ -446,12 +461,6 @@ class AM_IH(AM_IP):
        err_info = f" ({['EDC_FUE', 'ILLEGAL_INST', 'MEMVIOL', 'EDC_FED'][err_type]})" if enc_type == 2 else ""
        print(f"am {self.adev.devfmt}: sq_intr: {['auto', 'wave', 'error'][enc_type]}{err_info}")
        self.adev.is_err_state |= enc_type == 2
-      elif src_name == "UTCL2_FAULT" or (self.adev.ip_ver[am.GC_HWIP][0] == 9 and client == am.SOC15_IH_CLIENTID_UTCL2):
-        bf = self.adev.reg(self.adev.gmc.pf_status_reg('GC')).read_bitfields()
-        va = (self.adev.reg('regGCVM_L2_PROTECTION_FAULT_ADDR_HI32').read()<<32) | self.adev.reg('regGCVM_L2_PROTECTION_FAULT_ADDR_LO32').read()
-        print(f"am {self.adev.devfmt}: GCVM_L2_PROTECTION_FAULT_STATUS: {bf} {va<<12:#x}")
-        self.adev.reg('regGCVM_L2_PROTECTION_FAULT_CNTL').update(clear_protection_fault_status_addr=1)
-        self.adev.is_err_state = True
      else: self.adev.is_err_state = True

      rptr = (rptr + 8) % (self.ring_size // 4)
@@ -507,23 +516,21 @@ class AM_SDMA(AM_IP):
    for reg, inst in self.sdma_reginst:
      self.adev.reg(f"{reg}_RB_CNTL").update(rb_enable=0, inst=inst)
      self.adev.reg(f"{reg}_IB_CNTL").update(ib_enable=0, inst=inst)
-      self.adev.reg(f"{reg}_DOORBELL").update(enable=0, inst=inst)
-      self.adev.reg(f"{reg}_DOORBELL_OFFSET").update(offset=0, inst=inst)

    if self.adev.ip_ver[am.SDMA0_HWIP] >= (6,0,0):
      self.adev.regGRBM_SOFT_RESET.write(soft_reset_sdma0=1)
      time.sleep(0.01)
      self.adev.regGRBM_SOFT_RESET.write(0x0)

-  def setup_ring(self, ring_addr:int, ring_size:int, rptr_addr:int, wptr_addr:int, idx:int) -> int:
+  def setup_ring(self, ring_addr:int, ring_size:int, rptr_addr:int, wptr_addr:int, idx:int) -> tuple[int, int]:
    pipe, queue = idx // 4, idx % 4
    reg, inst = ("regSDMA_GFX", pipe+queue*4) if self.adev.ip_ver[am.SDMA0_HWIP][:2] == (4,4) else (f"regSDMA{pipe}_QUEUE{queue}", 0)
    doorbell = am.AMDGPU_NAVI10_DOORBELL_sDMA_ENGINE0 + (pipe+queue*4) * 0xA
    self.sdma_reginst.append((reg, inst))

    self.adev.reg(f"{reg}_MINOR_PTR_UPDATE").write(0x1, inst=inst)
-    self.adev.wreg_pair(f"{reg}_RB_RPTR", "", "_HI", 0, inst=inst)
-    self.adev.wreg_pair(f"{reg}_RB_WPTR", "", "_HI", 0, inst=inst)
+    if not self.adev.partial_boot: self.adev.wreg_pair(f"{reg}_RB_RPTR", "", "_HI", 0, inst=inst)
+    if not self.adev.partial_boot: self.adev.wreg_pair(f"{reg}_RB_WPTR", "", "_HI", 0, inst=inst)
    self.adev.wreg_pair(f"{reg}_RB_BASE", "", "_HI", ring_addr >> 8, inst=inst)
    self.adev.wreg_pair(f"{reg}_RB_RPTR_ADDR", "_LO", "_HI", rptr_addr, inst=inst)
    self.adev.wreg_pair(f"{reg}_RB_WPTR_POLL_ADDR", "_LO", "_HI", wptr_addr, inst=inst)
@@ -533,7 +540,7 @@ class AM_SDMA(AM_IP):
    self.adev.reg(f"{reg}_RB_CNTL").write(**({f'{self.sdma_name.lower()}_wptr_poll_enable':1} if self.adev.ip_ver[am.SDMA0_HWIP][:2]!=(4,4) else {}),
      rb_vmid=0, rptr_writeback_enable=1, rptr_writeback_timer=4, rb_enable=1, rb_priv=1, rb_size=(ring_size//4).bit_length()-1, inst=inst)
    self.adev.reg(f"{reg}_IB_CNTL").update(ib_enable=1, inst=inst)
-    return doorbell
+    return self.adev.reg(f"{reg}_RB_WPTR").read(inst=inst) | (self.adev.reg(f"{reg}_RB_WPTR_HI").read(inst=inst) << 32), doorbell

 class AM_PSP(AM_IP):
  def init_sw(self):
@@ -253,7 +253,7 @@ class HCQSignal(Generic[HCQDeviceType]):
    Raises RuntimeError if a fault is detected.
    """

-  def wait(self, value:int, timeout:int|None=None):
+  def wait(self, value:int, timeout:int=getenv("HCQDEV_WAIT_TIMEOUT_MS", 30000)):
    """
    Waits the signal is greater than or equal to a specific value.

@@ -261,7 +261,6 @@ class HCQSignal(Generic[HCQDeviceType]):
      value: The value to wait for.
      timeout: Maximum time to wait in milliseconds. Defaults to 30s.
    """
-    timeout = timeout or getenv("HCQDEV_WAIT_TIMEOUT_MS", 30000)
    start_time = int(time.perf_counter() * 1000)
    while (not_passed:=(prev_value:=self.value) < value) and (cur_time:=int(time.perf_counter() * 1000)) - start_time < timeout:
      self._sleep(cur_time - start_time)
@@ -326,7 +325,7 @@ class HCQProgram(Generic[HCQDeviceType]):
    return self.args_state_t(argsbuf, self, bufs, vals=vals)

  def __call__(self, *bufs:HCQBuffer, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1),
-               vals:tuple[int|None, ...]=(), wait:bool=False, timeout:int|None=None) -> float|None:
+               vals:tuple[int|None, ...]=(), wait:bool=False) -> float|None:
    """
    Enqueues the program for execution with the given arguments and dimensions.

@@ -350,7 +349,7 @@ class HCQProgram(Generic[HCQDeviceType]):

    q.signal(self.dev.timeline_signal, self.dev.next_timeline()).submit(self.dev)

-    if wait: self.dev.synchronize(timeout=timeout)
+    if wait: self.dev.synchronize()
    return (float(sig_en.timestamp - sig_st.timestamp) / 1e6) if wait else None

 class HCQCompiled(Compiled, Generic[SignalType]):
@@ -363,8 +362,7 @@ class HCQCompiled(Compiled, Generic[SignalType]):
  cpu_devices: list[HCQCompiled] = []

  def __init__(self, device:str, allocator:HCQAllocatorBase, compilers:CompilerSet, runtime, signal_t:Type[SignalType],
-               comp_queue_t:Callable[..., HWQueue], copy_queue_t:Callable[..., HWQueue]|None=None, kernargs_size=(16 << 20), sigalloc_size=0x1000,
-               can_recover:bool=False):
+               comp_queue_t:Callable[..., HWQueue], copy_queue_t:Callable[..., HWQueue]|None=None, kernargs_size=(16 << 20), sigalloc_size=0x1000):
    self.device_id:int = int(device.split(":")[1]) if ":" in device else 0

    from tinygrad.runtime.graph.hcq import HCQGraph
@@ -388,23 +386,22 @@ class HCQCompiled(Compiled, Generic[SignalType]):
    self.kernargs_buf:HCQBuffer = self.allocator.alloc(kernargs_size, BufferSpec(cpu_access=True))
    self.kernargs_offset_allocator:BumpAllocator = BumpAllocator(self.kernargs_buf.size, wrap=True)

-    self.can_recover = can_recover # Whether the device can recover from faults or timeouts
    self.error_state:Exception|None = None # Exception if error is unrecoverable and sync will always fail

    if self._is_cpu(): HCQCompiled.cpu_devices.append(self)

-  def synchronize(self, timeout:int|None=None):
+  def synchronize(self):
    if self.error_state is not None: raise self.error_state

    # If we have any work on CPU devices, need to synchronize them. This is just an optimization to release GIL allowing to finish faster.
    if not self._is_cpu():
      for dev in HCQCompiled.cpu_devices: dev.synchronize()

-    try: self.timeline_signal.wait(self.timeline_value - 1, timeout=timeout if timeout is not None and self.can_recover else None)
+    try: self.timeline_signal.wait(self.timeline_value - 1)
    except RuntimeError as e:
      self.error_state = e
      if hasattr(self, 'on_device_hang'): self.on_device_hang()
-      raise e
+      else: raise e

    if self.timeline_value > (1 << 31): self._wrap_timeline_signal()
    if PROFILE:
@@ -7,7 +7,7 @@ from tinygrad.uop.ops import consumer_map_from_toposort, gate_kernel_sink
 from tinygrad.uop.symbolic import symbolic, pm_simplify_valid, pm_drop_and_clauses
 from tinygrad.helpers import argsort, all_same, cpu_profile, PCONTIG, colored

-ALWAYS_CONTIGUOUS: set[Ops] = {Ops.CONTIGUOUS, Ops.ASSIGN, Ops.COPY, Ops.ALLREDUCE, Ops.BUFFER, Ops.BUFFER_VIEW,
+ALWAYS_CONTIGUOUS: set[Ops] = {Ops.CONTIGUOUS, Ops.ASSIGN, Ops.COPY, Ops.BUFFER, Ops.BUFFER_VIEW,
                     Ops.CONST, Ops.BIND, Ops.DEVICE, Ops.MSELECT, Ops.MSTACK, Ops.PARAM,
                     Ops.DEFINE_LOCAL, Ops.DEFINE_REG, Ops.LOAD, Ops.CALL, Ops.ENCDEC}

@@ -18,10 +18,6 @@ def realize_srcs(ctx:dict[UOp, None], rb:UOp) -> None:
    if s.base.op not in ALWAYS_CONTIGUOUS: ctx[s] = None

 def realize_assign_src(ctx:dict[UOp, None], buf:UOp, x:UOp):
-  # don't realize COPY/ALLREDUCE/BUFFER_VIEW/ENCDEC when they are the direct source of ASSIGN — the ASSIGN target buffer is the output
-  if x.op in {Ops.COPY, Ops.ALLREDUCE, Ops.BUFFER_VIEW, Ops.ENCDEC} and x in ctx \
-     and not buf.op_in_backward_slice_with_self(Ops.SHRINK, Ops.PERMUTE, Ops.FLIP, Ops.PAD):
-    del ctx[x]
  # you don't usually have to do this for assign unless there's a WAR hazard like TestAssign.test_assign_double_diamond_reduce
  if buf.base in x.backward_slice_with_self: ctx[x] = None

@@ -29,9 +25,9 @@ pm_generate_realize_map = PatternMatcher([
  # always realize SINK src
  (UPat(Ops.SINK, name="s"), lambda ctx,s: ctx.update((x.base, None) for x in s.src if x.base.op not in ALWAYS_CONTIGUOUS)),
  # always realize
-  (UPat({Ops.COPY, Ops.ALLREDUCE, Ops.BUFFER_VIEW, Ops.CONTIGUOUS, Ops.STORE, Ops.ASSIGN, Ops.ENCDEC}, name="tr"), realize),
+  (UPat({Ops.COPY, Ops.BUFFER_VIEW, Ops.CONTIGUOUS, Ops.STORE, Ops.ASSIGN, Ops.ENCDEC}, name="tr"), realize),
  # realize srcs of these
-  (UPat((Ops.COPY, Ops.ALLREDUCE, Ops.MSELECT, Ops.MSTACK, Ops.ENCDEC), name="rb"), realize_srcs),
+  (UPat((Ops.COPY, Ops.MSELECT, Ops.MSTACK, Ops.ENCDEC), name="rb"), realize_srcs),
  # sometimes realize src of assign
  (UPat(Ops.ASSIGN, src=(UPat.var("buf"), UPat.var("x"))), realize_assign_src),
 ])
@@ -56,7 +52,7 @@ class IndexingContext:
    return UOp.range(s, next(self.range_idx), axistype) if resolve(s!=1) else UOp.const(dtypes.index, 0)

 def create_bufferize_and_index_based_on_ranges(ctx:IndexingContext, x:UOp):
-  if x.op in {Ops.BUFFERIZE, Ops.INDEX}: return None
+  if x.op in {Ops.BUFFERIZE, Ops.INDEX, Ops.AFTER}: return None
  new_srcs = []
  for s in x.src:
    new_src = s
@@ -71,8 +67,8 @@ def create_bufferize_and_index_based_on_ranges(ctx:IndexingContext, x:UOp):
        new_src = s.end(*[r for r in closed_ranges if r.op is Ops.RANGE])
        del ctx.realize_map[s]
      else:
-        # the Bufferize before a COPY/ALLREDUCE is not removable. there should be a better way to do this
-        removable = x.op not in {Ops.COPY, Ops.ALLREDUCE} and s.op not in ALWAYS_CONTIGUOUS
+        # the Bufferize before a COPY is not removable. there should be a better way to do this
+        removable = x.op is not Ops.COPY and s.op not in ALWAYS_CONTIGUOUS
        # None in the device assigns it a number later
        opts = BufferizeOpts(device=s.device, removable=removable) if len(ctx.range_map[s][1]) == len(realized_ranges) else \
               BufferizeOpts(device=s.device, addrspace=AddrSpace.LOCAL, removable=removable)
@@ -122,6 +118,8 @@ pm_apply_rangeify = PatternMatcher([
  (UPat(GroupOp.All, name="x"), create_bufferize_and_index_based_on_ranges),
  # remove movement op
  (UPat(GroupOp.Movement, name="x"), remove_movement_op_after_rangeify),
+  # const/define_var shouldn't have src
+  (UPat((Ops.CONST, Ops.DEFINE_VAR), name="c"), lambda ctx,c: c.replace(src=()) if c in ctx.range_map else None),
 ])

@functools.cache
@@ -148,12 +146,13 @@ def apply_movement_op(op:Ops, in_shape:tuple[sint,...], arg:tuple, rngs:tuple[UO
    case Ops.FLIP:    rngs = tuple(((s-1)-a) if f else a for a,s,f in zip(rngs, in_shape, arg))
    case Ops.EXPAND:  rngs = tuple(a if in_sh == out_sh else a.const_like(0) for a,in_sh,out_sh in zip(rngs, in_shape, arg))
    case Ops.PAD:
-      # NOTE: the .where(r-s, i) is not inside the graph_rewrite so that `convert_pad_to_where_to_keep_behavior_local`
+      # TODO: why is multiple graph_rewrites faster than one here?
+      # TODO: the .where(r-s, i) is not inside the graph_rewrite so that `convert_pad_to_where_to_keep_behavior_local`
      #       wraps the pad with only the newly added valid
-      rngs = tuple(r if (s == 0 and e == 0) else graph_rewrite((r >= s) & (r < (sh+s)),
+      rngs = tuple(r if (s == 0 and e == 0) else graph_rewrite(((r >= s) & (r < (sh+s))),
        symbolic+pm_simplify_valid, name="pad").where(r-s, UOp.invalid()) for r,sh,(s,e) in zip(rngs, in_shape, arg))
    case Ops.RESHAPE:
-      sink = UOp.sink(*rngs).simplify() # NOTE: this applies any commutative flips to the rngs early
+      sink = UOp.sink(*rngs)
      sub_array = {r:UOp.range(r.src[0], i, AxisType.PLACEHOLDER) for i,r in enumerate(sink.ranges)}
      rngs = _apply_reshape(in_shape, arg, sink.substitute(sub_array)).substitute({v:k for k,v in sub_array.items()}).src
    case _: raise RuntimeError(f"{op} is not a MovementOp")
@@ -165,7 +164,12 @@ def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]:
  rctx = IndexingContext()

  # get ops to realize
-  graph_rewrite(tsink, pm_generate_realize_map, ctx=rctx.realize_map, name="get realize")
+  graph_rewrite(tsink, pm_generate_realize_map, ctx=rctx.realize_map, bottom_up=True, name="get realize")
+  # don't realize COPY/BUFFER_VIEW/ENCDEC when they are the direct source of ASSIGN — the ASSIGN target buffer is the output
+  for u in tsink.toposort():
+    if u.op is Ops.ASSIGN and u.src[1].op in {Ops.COPY, Ops.BUFFER_VIEW, Ops.ENCDEC} and u.src[1] in rctx.realize_map \
+       and not u.src[0].op_in_backward_slice_with_self(Ops.SHRINK, Ops.PERMUTE, Ops.FLIP, Ops.PAD):
+      del rctx.realize_map[u.src[1]]

  # get the consumer map
  with cpu_profile("consumer map in rangeify", "TINY"):
@@ -177,13 +181,7 @@ def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]:
    if x.op in {Ops.DEVICE, Ops.UNIQUE}: continue

    # no ranges on kernels, they are internal
-    if x.op in {Ops.CALL, Ops.LINEAR}: continue
-
-    # no range on after
-    if x.op is Ops.AFTER: continue
-
-    # treat MSTACK/MSELECT like SINK
-    if x.op in {Ops.MSTACK, Ops.MSELECT}: continue
+    if x.op is Ops.CALL: continue

    if x.dtype.scalar() == dtypes.index: continue  # TODO: why do I need this?
    ending_ranges[x] = sum([ending_ranges.get(u, []) for u in consumer_map[x]], [])
@@ -202,6 +200,9 @@ def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]:
      # mark all ranges as ended
      assert rctx.realize_map[x] is None
      rctx.realize_map[x] = list(range(len(x.shape)))
+    elif x.op in {Ops.MSTACK, Ops.MSELECT}:
+      # treat MSTACK/MSELECT like SINK
+      continue
    elif len(consumer_rngs) == 0:
      # if no consumers have ranges and this isn't realized, this doesn't have ranges either.
      continue
@@ -236,7 +237,7 @@ def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]:

    # if this element is a reduce and there's ended ranges, we might have to end some other ranges
    if len(ending_ranges[x]) and x.op in GroupOp.Elementwise.union({Ops.REDUCE_AXIS}):
-      _realize_axis = rctx.realize_map.get(x) or []
+      _realize_axis = rctx.realize_map.get(x, []) or []
      for i,r in enumerate(out_rngs):
        if i in _realize_axis: continue
        if not (PCONTIG > 1) or any(any(rr.arg > e.arg for e in ending_ranges[x]) for rr in r.ranges):
@@ -1,6 +1,6 @@
 import functools, itertools
 from tinygrad.helpers import all_same, all_int, prod, DEBUG, RING, ALL2ALL, getenv
-from tinygrad.uop.ops import Ops, UOp, PatternMatcher, UPat, GroupOp, graph_rewrite, should_resolve_call
+from tinygrad.uop.ops import Ops, UOp, PatternMatcher, UPat, GroupOp
 from tinygrad.dtype import dtypes

 # *** allreduce implementation ***
@@ -71,7 +71,7 @@ def mstack_early_shrink(ms:UOp, shrink:UOp):
  return ms.replace(src=tuple(ret))

 replace_allreduce = PatternMatcher([
-  #(UPat(Ops.ALLREDUCE, src=(UPat.var("buf"), UPat()), name="red"), handle_allreduce),
+  (UPat(Ops.ALLREDUCE, src=(UPat.var("buf"), UPat()), name="red"), handle_allreduce),
  # BROADCAST: explicitly expand broadcast copies and combine with MSTACK
  (UPat(Ops.COPY, name="c", src=(UPat(GroupOp.All-{Ops.CONST}, name="x"), UPat(Ops.DEVICE))), lambda c,x:
    UOp(Ops.MSTACK, c.dtype, tuple(x.copy_to_device(d) for d in c.device)) if isinstance(c.device, tuple) and isinstance(x.device, str) else None),
@@ -163,19 +163,8 @@ def assign_multi(dest:UOp, src:UOp):
 def passthrough_multi(root:UOp, multi:UOp):
  return UOp(root.op, root.dtype, (multi.src[0],)+tuple(x.src[0] if x.op is Ops.MULTI else x for x in root.src[1:]), root.arg).multi(multi.axis)

-def rewrite_into_call(call:UOp):
-  if not should_resolve_call(call): return None
-  new_body = graph_rewrite(call.src[0], multi_pm, name="subcall")
-  new_args = tuple(a.src[0] if a.op is Ops.MULTI else a for a in call.src[1:])
-  return call.replace(src=(new_body,)+new_args)
-
-def param_to_multi(p:UOp):
-  if p.axis is None: return None
-  return UOp.param(p.arg, p.dtype, p.shard_shape, p._device).multi(p.axis)
-
 # NOTE: this is the same pattern as Ops.UNROLL
 multi_pm = PatternMatcher([
-  (UPat(Ops.PARAM, name="p"), param_to_multi),
  (UPat(GroupOp.ALU, name="root", custom_early_reject=set([Ops.MULTI])), alu_multi),
  (UPat(Ops.REDUCE_AXIS, src=(UPat(Ops.MULTI, name="multi"), ), name="root"), reduce_multi),
  (UPat(Ops.RESHAPE, src=(UPat(Ops.MULTI, name="multi"), UPat()), name="root"), reshape_multi),
@@ -188,8 +177,6 @@ multi_pm = PatternMatcher([
  (UPat(Ops.COPY, src=(UPat(Ops.MULTI, name="multi"), UPat(Ops.DEVICE, name="device"))), copy_multi),
  (UPat(Ops.ALLREDUCE, src=(UPat(Ops.MULTI, name="multi"), UPat(Ops.DEVICE, name="device")), name="red"),
    lambda multi,device,red: multi.src[0].allreduce(red.arg, device).multi(axis=multi.axis)),
-  # rewrite into calls explicitly for MULTI
-  (UPat(Ops.CALL, name="call"), rewrite_into_call),
  (UPat(Ops.CALL, src=(UPat(Ops.MULTI, name="multi"), ), name="root", allow_any_len=True), passthrough_multi),
  # we just remove the MULTI from CALLs with dtypes.void and assume they are handled by the user for custom kernels
  (UPat(Ops.CALL, dtype=dtypes.void, name="root", custom_early_reject=set([Ops.MULTI])), lambda root:
@@ -1,8 +1,8 @@
 from dataclasses import dataclass, field, replace
 import itertools
-from tinygrad.dtype import dtypes, PtrDType, ImageDType, AddrSpace, Invalid
+from tinygrad.dtype import dtypes, PtrDType, ImageDType, AddrSpace
 from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, _substitute, KernelInfo
-from tinygrad.uop.ops import graph_rewrite, sint, AxisType, BottomUpGate, profile_matches, should_resolve_call, identity_element
+from tinygrad.uop.ops import graph_rewrite, sint, AxisType, BottomUpGate, profile_matches
 from tinygrad.uop.symbolic import symbolic
 from tinygrad.helpers import prod, all_same, getenv, dedup, all_int, DEBUG, SPLIT_REDUCEOP, DEBUG_RANGEIFY, VIZ, MAX_KERNEL_BUFFERS
 from tinygrad.helpers import PCONTIG, partition, get_single_element
@@ -76,39 +76,30 @@ mop_cleanup = PatternMatcher([
 ])

 pm_gather_params = PatternMatcher([ (UPat(Ops.PARAM, name="p"), lambda ctx, p: ctx.append(p)), ])
-def resolve_call(c:UOp, allow_param_mismatch=True) -> UOp|None:
-  if not should_resolve_call(c): return None
+def resolve_call(c:UOp, allow_param_mismatch=False) -> UOp|None:
+  # don't resolve real kernel calls, sink or program
+  if c.src[0].op is Ops.SINK and isinstance(c.src[0].arg, KernelInfo): return None
+  if c.src[0].op is Ops.PROGRAM: return None
  params: list[UOp] = []
-  graph_rewrite(c.src[0], pm_gather_params, bottom_up=True, ctx=params, name="gather params")
+  graph_rewrite(c.src[0], pm_gather_params, bottom_up=True, ctx=params)
  params = sorted(params, key=lambda x: x.arg)
  args = c.src[1:]
-
-  # NOTE: this isn't really needed. it's okay if there's unused args in the function
+  # TODO: this check belongs in spec, not here
  if not allow_param_mismatch:
    if [x.arg for x in params] != list(range(len(params))): raise RuntimeError(f"params not in order: {[x.arg for x in params]}")
    if len(params) != len(args): raise TypeError(f"expected {len(params)} args, got {len(args)}")
-
-  dict_map = {x:args[x.arg] for x in params}
-  for i, (p, a) in enumerate(dict_map.items()):
-    if p.axis != a.axis: raise TypeError(f"arg {i} axis mismatch: expected {p.axis}, got {a.axis}")
-    if p.max_shape != a.max_shape: raise TypeError(f"arg {i} shape mismatch: expected {p.shape}, got {a.shape}")
+  for i, (p, a) in enumerate(zip(params, args)):
+    if p.shape != a.shape: raise TypeError(f"arg {i} shape mismatch: expected {p.shape}, got {a.shape}")
    if p.dtype != a.dtype: raise TypeError(f"arg {i} dtype mismatch: expected {p.dtype}, got {a.dtype}")
-  return c.src[0].substitute(dict_map, walk=True)
+  return c.src[0].substitute(dict(zip(params, args)))

 earliest_rewrites = mop_cleanup+PatternMatcher([
-  # early fixup const copy
-  (UPat(Ops.COPY, src=(UPat.var("s"), UPat.var("d"))),
-   lambda s,d: s.substitute({UOp(Ops.DEVICE, arg=s.device):d}) if s.base.op is Ops.CONST else None),
-
  # resolve calls
  (UPat(Ops.CALL, name="c"), resolve_call),

  # split_reduceop
  (UPat(Ops.REDUCE_AXIS, name="reduce", src=(UPat.var("x"),)), split_reduceop),

-  # remove DETACH/CONTIGUOUS_BACKWARD (TODO: this is copied in allocations)
-  (UPat((Ops.DETACH, Ops.CONTIGUOUS_BACKWARD), name="x"), lambda x: x.src[0]),
-
  # remove contiguous on movement ops before a copy on disk
  (UPat(GroupOp.Movement-{Ops.SHRINK, Ops.RESHAPE}, name="x").f(Ops.CONTIGUOUS).f(Ops.COPY, allow_any_len=True, name="copy"),
   lambda x,copy: copy.replace(src=(x,)+copy.src[1:]) if isinstance(x.device, str) and x.device.startswith("DISK") else None),
@@ -119,8 +110,8 @@ earliest_rewrites = mop_cleanup+PatternMatcher([

  # ** copy rules **

-  # COPY/ALLREDUCE and source size need to match
-  (UPat((Ops.COPY, Ops.ALLREDUCE), src=(UPat(GroupOp.Movement, name="r"), UPat(name="d")), name="c"),
+  # COPY and source size need to match
+  (UPat(Ops.COPY, src=(UPat(GroupOp.Movement, name="r"), UPat(name="d")), name="c"),
   lambda c,r,d: c.replace(src=(r.contiguous(), d)) if r.size != r.base.size else None),

  # copy only to different device
@@ -140,20 +131,12 @@ earliest_rewrites = mop_cleanup+PatternMatcher([

  # make source contiguous if it has hazardous movement ops on the dest buffer
  (UPat(Ops.ASSIGN, src=(UPat.var("target"), UPat.var("src")), name="assign"), fix_assign_hazard),
-
-  # ** size 0 **
-
-  # reduce of size 0 is the identity element
-  (UPat(Ops.REDUCE_AXIS, name="reduce", src=(UPat.var("x"),)),
-   lambda reduce,x: reduce.const_like(identity_element(reduce.arg[0], reduce.dtype)) if x.size == 0 and reduce.size != 0 else None),
-  # handle size 0
-  (UPat(GroupOp.All-{Ops.SINK}, name="x"), lambda x: x.const_like(0).rtag(x.tag) if x._shape is not None and x.size == 0 else None),
 ])

 # *****************
 # 3.5 cleanups

-ALWAYS_RUN_OPS = {Ops.CONTIGUOUS, Ops.COPY, Ops.ALLREDUCE, Ops.ASSIGN, Ops.ENCDEC, Ops.NOOP}
+ALWAYS_RUN_OPS = {Ops.CONTIGUOUS, Ops.COPY, Ops.ASSIGN, Ops.ENCDEC, Ops.NOOP}

 # you don't know in the first pass if axes are going to die, this happens if there's an EXPAND to the left
 def cleanup_dead_axes(b:UOp):
@@ -242,9 +225,8 @@ def remove_bufferize(src:UOp, buf:UOp, idx:UOp):

  # if it makes it here, the bufferize is removed
  # this is the ranges replaced
-  # NOTE: if buf src is a const, we don't replace it. if idx is Invalid (dead load), don't replace it either
-  replaced = {k:v for k,v in zip(buf.src[1:], idx.src[1:]) if k.op is not Ops.CONST and not (v.op is Ops.CONST and v.arg is Invalid)}
-  return src.substitute(replaced, extra_pm=pm_gate_substitute)
+  # NOTE: if buf src is a const, we don't replace it
+  return src.substitute({k:v for k,v in zip(buf.src[1:], idx.src[1:]) if k.op is not Ops.CONST}, extra_pm=pm_gate_substitute)

 def remove_noop_bufferize(idx,b2):
  if idx.src[1:] != b2.src[1:] or idx.src[0].op is Ops.BUFFER_VIEW: return None
@@ -263,8 +245,6 @@ pm_const_buffer_folding = pm_mops+PatternMatcher([
  (UPat(Ops.INDEX, src=(UPat(Ops.CONST, name="c"),),), lambda c: c),
  # copy on CONST is CONST
  (UPat(Ops.COPY, src=(UPat.cvar("x"), UPat()), name="copy"), lambda copy,x: copy.const_like(x.arg)),
-  # allreduce on CONST is CONST
-  (UPat(Ops.ALLREDUCE, src=(UPat.cvar("x"), UPat()), name="copy", arg=Ops.ADD), lambda copy,x: copy.const_like(x.arg)*len(x.device)),
  # hack if a noop turned to a const
  (UPat(Ops.NOOP, src=(UPat.cvar("c"),), name="noop"), lambda c,noop: c),
  # mstack on CONST is CONST
@@ -380,11 +360,6 @@ pm_add_buffers = pm_mops+pm_flatten_bufferize+to_bufferview+PatternMatcher([

  # remove any RESHAPEs on KERNEL
  (UPat(Ops.CALL, name="k"), lambda k: k.replace(src=tuple(x.src[0] if x.op is Ops.RESHAPE else x for x in k.src))),
-
-  # remove MOP on AFTER
-  (UPat(Ops.AFTER, src=(UPat.var("x"), UPat(GroupOp.Movement, name="y"))), lambda x,y: x.after(y.src[0])),
-  # remove double AFTER
-  (UPat(Ops.AFTER, src=(UPat.var("x"), UPat(Ops.AFTER, name="y"))), lambda x,y: x.after(*y.src[1:]))
 ])

 pm_add_buffers_local = pm_mops+pm_flatten_bufferize+to_bufferview+PatternMatcher([
@@ -403,7 +378,7 @@ class LocalAddBufferContext:
  opts:tuple|None = None

 def debuf(ctx:LocalAddBufferContext, buf:UOp):
-  ret = UOp(Ops.PARAM, buf.dtype.ptr(buf.size), arg=ctx.dg).reshape(buf.shape)
+  ret = UOp(Ops.PARAM, buf.dtype.ptr(buf.size), arg=ctx.dg)
  if buf not in ctx.map: ctx.map[buf] = buf
  ctx.dg += 1
  return ret
@@ -492,8 +467,7 @@ def split_store(x:UOp) -> UOp|None:
  if ret.op is Ops.STORE: stored = ret.src[1]
  elif ret.op is Ops.END and ret.src[0].op is Ops.STORE: stored = ret.src[0].src[1]
  else: raise RuntimeError(f"unknown kernel type {ret.op}")
-  if stored.op in {Ops.COPY, Ops.ALLREDUCE, Ops.BUFFER_VIEW}: ret = stored.replace(src=stored.src + ret.ended_ranges)
-  elif stored.op is Ops.ENCDEC: ret = stored
+  if stored.op in {Ops.COPY, Ops.BUFFER_VIEW, Ops.ENCDEC}: ret = stored
  else: ret = ret.sink(arg=KernelInfo(opts_to_apply=lctx.opts))

  kernel = ret.call(*lctx.map.values(), *lctx.vars.keys())
@@ -507,7 +481,7 @@ split_kernels = PatternMatcher([

@profile_matches
 def get_kernel_graph(sink:UOp) -> UOp:
-  tsink = graph_rewrite(sink, multi_pm, name="multi_pm")
+  tsink = graph_rewrite(sink, multi_pm, name="multi_pm", rewrite_into_calls=True)
  tsink = graph_rewrite(tsink, pm_syntactic_sugar+pm_mops+earliest_rewrites, bottom_up=True, name="earliest rewrites")

  # convert movement ops to ranges
@@ -537,4 +511,4 @@ def get_kernel_graph(sink:UOp) -> UOp:
      assign_rep[a] = kernel_assign[s] = a.replace(src=a.src+(u,))
  if assign_rep: tsink = graph_rewrite(tsink, _substitute, ctx=assign_rep, bottom_up=True, name="fix_assign")
  if VIZ: graph_rewrite(tsink, PatternMatcher([]), name="View Kernel Graph")
-  return tsink
+  return tsink
--- a/Show More
+++ b/Show More
Author	SHA1	Message	Date
geohot	8395071f77	recursive stuff works	2026-02-24 15:15:36 +08:00
geohot	de3e901b71	works but bad	2026-02-24 14:40:39 +08:00
geohot	ae2410e10e	add callify method	2026-02-24 11:44:33 +08:00