Merge branch 'master' into sym_work

.github/actions/setup-tinygrad/action.yml

@@ -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

.github/workflows/benchmark.yml

@@ -520,9 +520,8 @@ 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
-    #- 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: 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
       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
 
@@ -617,27 +616,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]

.github/workflows/test.yml

@@ -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)

examples/mlperf/dataloader.py

@@ -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}")

examples/mlperf/model_train.py

@@ -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
 
@@ -1336,12 +1336,10 @@ def train_llama3():
   # 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
+  if (llama_layers:=getenv("LLAMA_LAYERS")) != 0: model_params['n_layers'] = llama_layers
+  print(f"model parameters: {model_params}")
 
   model = Transformer(**model_params, max_context=SEQLEN, jit=False, disable_kv_cache=True)
   params = get_parameters(model)
@@ -1387,10 +1385,8 @@ def train_llama3():
 
   # 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 +1401,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 = vocab_mask.where(-float("inf"), 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 +1417,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 = vocab_mask.where(-float("inf"), 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):
@@ -1484,7 +1478,7 @@ def train_llama3():
 
       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:
@@ -1497,8 +1491,7 @@ def train_llama3():
       if stopped: break
 
       gt = time.perf_counter()
-      ret = optim_step()
-      lr, grad_norm = ret[0].item(), ret[1].item()
+      lr = optim_step().item()
       et = time.perf_counter()
 
       loss = sum(losses) / len(losses)
@@ -1516,14 +1509,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 +1550,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()

examples/mlperf/optim.py

@@ -7,49 +7,40 @@ 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:
       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()
+        grads[i] = grads[i] / self.grad_acc
+      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.clip_norm / (total_norm + 1e-6)).clamp(max_=1.0)).cast(grads[i].dtype)).realize()
+        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()

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

@@ -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}

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

@@ -3,4 +3,4 @@ export BENCHMARK=5
 export EVAL_BS=0
 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

examples/openpilot/load_pickle.py

@@ -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"

extra/bench_log.py

@@ -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:

extra/gemm/asm/cdna/gemm.py

@@ -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
@@ -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,16 +85,9 @@ 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)
 
@@ -117,6 +98,4 @@ def asm_gemm(a:Tensor, b:Tensor) -> Tensor:
   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

extra/sqtt/README.md

@@ -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
+```

extra/sqtt/active_sqtt_parse.py

@@ -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)

extra/sqtt/attempt_sqtt_parse.py

@@ -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))

extra/sqtt/examples/discover_ops.py

@@ -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)

extra/sqtt/examples/generate_examples.py

@@ -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__":

extra/sqtt/roc.py

@@ -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()

extra/torch_backend/test_kernel_fusion.py

@@ -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()

extra/usbgpu/tbgpu/installer/TinyGPUDriverExtension/TinyGPUDriver.NV.Release.entitlements

@@ -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>

extra/usbgpu/tbgpu/installer/build_and_sign_nv.sh

@@ -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

extra/viz/README

@@ -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"`

extra/viz/cli.py

@@ -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

test/amd/disasm.py

@@ -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}"

test/amd/hw/test_vopc.py

@@ -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 = [

test/amd/test_llvm.py

@@ -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]]:

test/amd/test_sqtt_examples.py

@@ -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

test/amd/test_sqtt_tables.py

@@ -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()

test/amd/test_sqttmap.py

@@ -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"

test/backend/test_asm_gemm.py

@@ -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)

test/backend/test_const_folding.py

@@ -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):

test/backend/test_multitensor.py

@@ -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):

test/backend/test_schedule.py

@@ -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))

test/backend/test_setitem.py

@@ -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):

test/backend/test_stunning.py

@@ -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)
 

test/backend/test_subbuffer.py

@@ -1,7 +1,7 @@
 import unittest
 from tinygrad import Device, dtypes, Tensor
 from tinygrad.device import Buffer
-from tinygrad.helpers import Context, getenv
+from tinygrad.helpers import Context
 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()

test/backend/test_symbolic_ops.py

@@ -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)
 

test/backend/test_tensor_variable.py

@@ -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()

test/external/external_fuzz_beam_timeout_recovery.py

@@ -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 ===")

test/external/external_fuzz_sdma_warm_start.py

@@ -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 ===")

test/external/external_test_gpu_crash.py

@@ -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))
 
 

test/external/external_test_speed_llama.py

@@ -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

test/mockgpu/amd/pcode.py

@@ -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 '+' | '-':

test/models/test_onnx.py

@@ -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()

test/null/test_mnist_dataset.py

@@ -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()

test/null/test_multitensor.py

@@ -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 = {}

test/null/test_schedule.py

@@ -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)

test/null/test_viz.py

@@ -1,4 +1,4 @@
-import unittest, decimal, sys
+import unittest, decimal, json, struct, sys
 from dataclasses import dataclass
 from typing import Generator
 
@@ -357,9 +357,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):

test/unit/test_assign.py

@@ -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()
@@ -609,8 +608,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 +686,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 +766,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.

test/unit/test_call.py

@@ -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()

test/unit/test_function.py

@@ -70,14 +70,6 @@ class TestFunction(unittest.TestCase):
     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])
@@ -137,31 +129,6 @@ class TestFunction(unittest.TestCase):
     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
@@ -226,114 +193,5 @@ class TestFunction(unittest.TestCase):
     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()

test/unit/test_gguf.py

@@ -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()

test/unit/test_gradient.py

@@ -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)

test/unit/test_indexing.py

@@ -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)

test/unit/test_linalg.py

@@ -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()

tinygrad/apps/llm.py

@@ -143,14 +143,13 @@ class TransformerBlock:
     #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)))
+  @function
   def _feed_forward(self, h: Tensor) -> Tensor:
     h_norm = self.ffn_norm(h)
     if hasattr(self, 'ffn_gate_exps'):
@@ -190,9 +189,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 +219,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 +262,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 +336,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)

tinygrad/codegen/__init__.py

@@ -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)
 

tinygrad/codegen/late/devectorizer.py

@@ -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([

tinygrad/codegen/opt/search.py

@@ -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

tinygrad/device.py

@@ -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

tinygrad/engine/allocations.py

@@ -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, track_rewrites
+from tinygrad.dtype import ImageDType
+from tinygrad.helpers import prod, DEBUG, argsort, VIZ, pluralize
 
 @dataclass
 class AllocCtx:
@@ -25,22 +25,24 @@ def disk_copy_is_buffer(ctx:AllocCtx, u:UOp):
   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/TINYFS tensors, they are left alone
+  if isinstance(u._device, str) and u._device.startswith(("DISK", "TINYFS")): 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):
@@ -160,8 +121,6 @@ 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),
 ])
@@ -183,6 +142,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

tinygrad/engine/realize.py

@@ -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)

tinygrad/engine/schedule.py

@@ -2,9 +2,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 _remove_all_tags
 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 +23,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, Ops.LINEAR}, f"AFTER src[1] should be CALL 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:]
@@ -91,31 +92,27 @@ def create_new_buffer(ctx:tuple[dict[UOp, UOp], tuple[UOp, ...]], b:UOp):
   return ret
 
 pm_post_sched_cache = PatternMatcher([
-  (UPat(Ops.PARAM, name="x"), lambda ctx,x: ctx[1][x.arg]),
+  (UPat(Ops.PARAM, name="x"), lambda ctx,x: ctx[1][x.arg].rtag() 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),
 ])
 
-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)))),
+# the AFTER structure is already in LINEAR
+pm_collapse_after = PatternMatcher([
+  (UPat(Ops.AFTER, name="x"), lambda x: x.src[0])
 ])
 
 schedule_cache: dict[bytes, UOp] = {}
-# ctx is just for DEBUG on inner
-def lower_sink_to_linear(function:UOp) -> UOp|None:
+def lower_schedule_to_linear(big_sink:UOp) -> UOp|None:
   st = time.perf_counter()
+  function = big_sink.src[0]
   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)
+  if not SCACHE or (sc_ret:=schedule_cache.get(function.key, None)) is None:
+    if SPEC: type_verify(big_sink, tensor_spec)
     # support recursive CALLs
+    function = graph_rewrite(function, pm_schedule, name="inner schedule to linear")
     linear = create_schedule(get_kernel_graph(function))
-    if SCACHE: schedule_cache[cache_key] = linear
+    if SCACHE: schedule_cache[function.key] = linear
   else:
     # schedule cache hit
     linear = sc_ret
@@ -127,30 +124,20 @@ def lower_sink_to_linear(function:UOp) -> UOp|None:
     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]}"+\
+          f" | {' cache hit' if SCACHE and sc_ret is not None else 'CACHE MISS'} {function.key.hex()[:8]}"+\
           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:])
+  # TODO: use walk and avoid the remove tags
+  linear = graph_rewrite(linear, pm_post_sched_cache, ctx=({}, big_sink.src[1:]), walk=True, name="params to buffers")
+  return graph_rewrite(linear, pm_collapse_after+_remove_all_tags, name="remove tags/after")
 
 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),
+  (UPat(Ops.CALL, src=(UPat(Ops.SINK),), allow_any_len=True, name="big_sink"), lower_schedule_to_linear),
 ])
 
 @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")
+  linear = graph_rewrite(big_sink, pm_schedule, name="schedule to linear")
 
   # vars used in the schedule
   used_vars = set().union(*[{v.expr for v in si.src[0].variables()} for si in linear.src])
@@ -161,7 +148,7 @@ def complete_create_schedule_with_vars(big_sink:UOp) -> tuple[list[ExecItem], di
       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}")
+      assert nm not in var_vals or var_vals[nm] == val, f"bind mismatch on {nm}, {var_vals[nm]} != {val}"
       var_vals[nm] = val
 
   # convert LINEAR to ExecItems

tinygrad/function.py

@@ -1,5 +1,5 @@
 import functools
-from typing import Generic, TypeVar, Callable, cast, overload
+from typing import Generic, TypeVar, Callable, cast
 from tinygrad.helpers import Context, dedup, getenv
 from tinygrad.uop.ops import UOp, Ops, graph_rewrite, PatternMatcher, UPat
 from tinygrad.tensor import Tensor
@@ -16,10 +16,9 @@ pm_ctx = PatternMatcher([
 ])
 
 ReturnType = TypeVar('ReturnType')
-class _function(Generic[ReturnType]):
-  def __init__(self, fxn:Callable[..., ReturnType], *, precompile:bool=False):
+class function(Generic[ReturnType]):
+  def __init__(self, fxn:Callable[..., ReturnType]):
     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
 
@@ -58,14 +57,6 @@ class _function(Generic[ReturnType]):
     #call = assigned.call(*call_uops, buffer, name=name)
     #ret = buffer.after(call)
 
-    ret = uret.call(*call_uops, name=name, precompile=self.precompile)
+    ret = uret.call(*call_uops, name=name)
     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)

tinygrad/gradient.py

@@ -17,13 +17,12 @@ def call_gradient(ctx:UOp, k:UOp) -> tuple[UOp|None, ...]:
   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()))
+  params = sorted([x for x in fxn.toposort() if x.op == Ops.PARAM], key=lambda x: x.arg)
+  grads = compute_gradient(fxn, ctx.param_like(len(args)), set(params))
   ret: list[UOp|None] = [None]
-  for i in range(len(args)):
-    if (p:=params.get(i, None)) is not None and p in grads:
+  for i,p in enumerate(params):
+    if 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)

tinygrad/helpers.py

@@ -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
 

tinygrad/nn/optim.py

@@ -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):
     """

tinygrad/nn/state.py

@@ -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 ])

tinygrad/renderer/amd/sqtt.py

@@ -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))

tinygrad/runtime/ops_amd.py

@@ -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,21 +857,21 @@ 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)
+      if reset and 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
 
@@ -978,8 +977,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

tinygrad/runtime/ops_cl.py

@@ -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]:

tinygrad/runtime/ops_cuda.py

@@ -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)

tinygrad/runtime/ops_dsp.py

@@ -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()

tinygrad/runtime/ops_hip.py

@@ -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))]

tinygrad/runtime/ops_metal.py

@@ -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()

tinygrad/runtime/ops_null.py

@@ -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']):

tinygrad/runtime/ops_nv.py

@@ -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():

tinygrad/runtime/ops_python.py

@@ -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)

tinygrad/runtime/ops_qcom.py

@@ -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=}")

tinygrad/runtime/ops_webgpu.py

@@ -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

tinygrad/runtime/support/am/amdev.py

@@ -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,14 +243,14 @@ 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):

tinygrad/runtime/support/am/ip.py

@@ -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)}:
@@ -240,11 +240,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 +285,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 +320,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 +390,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 +445,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}]")
 
@@ -450,7 +459,6 @@ class AM_IH(AM_IP):
         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
 
@@ -507,23 +515,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 +539,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):

tinygrad/runtime/support/hcq.py

@@ -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:

tinygrad/schedule/indexing.py

@@ -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,8 +18,8 @@ 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 \
+  # don't realize COPY/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.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
@@ -29,9 +29,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),
 ])
@@ -71,8 +71,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)
@@ -153,7 +153,7 @@ def apply_movement_op(op:Ops, in_shape:tuple[sint,...], arg:tuple, rngs:tuple[UO
       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")

tinygrad/schedule/multi.py

@@ -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),
@@ -164,18 +164,10 @@ 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)
+  return call.replace(src=(graph_rewrite(call.src[0], multi_pm, name="subcall"),)+call.src[1:]) if should_resolve_call(call) else None
 
 # 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),

tinygrad/schedule/rangeify.py

@@ -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, should_resolve_call
 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
@@ -90,16 +90,11 @@ def resolve_call(c:UOp, allow_param_mismatch=True) -> UOp|None:
 
   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}")
     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)
 
 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),
 
@@ -119,8 +114,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 +135,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 +229,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 +249,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
@@ -403,7 +387,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,7 +476,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)
+  if stored.op in {Ops.COPY, Ops.BUFFER_VIEW}: ret = stored.replace(src=stored.src + ret.ended_ranges)
   elif stored.op is Ops.ENCDEC: ret = stored
   else: ret = ret.sink(arg=KernelInfo(opts_to_apply=lctx.opts))
 

tinygrad/tensor.py

@@ -7,7 +7,7 @@ if TYPE_CHECKING: import numpy
 from tinygrad.dtype import DType, DTypeLike, dtypes, ImageDType, ConstType, least_upper_float, least_upper_dtype, sum_acc_dtype, to_dtype, truncate
 from tinygrad.dtype import _from_np_dtype, _to_np_dtype, PyConst
 from tinygrad.helpers import argfix, make_tuple, flatten, prod, all_int, round_up, merge_dicts, argsort, getenv, all_same, fully_flatten
-from tinygrad.helpers import IMAGE, FLOAT16, WINO, Metadata, TRACEMETA, ASM_GEMM, ceildiv, fetch, is_numpy_ndarray, TracingKey, cpu_profile
+from tinygrad.helpers import IMAGE, WINO, Metadata, TRACEMETA, ASM_GEMM, ceildiv, fetch, is_numpy_ndarray, TracingKey, cpu_profile
 from tinygrad.helpers import suppress_finalizing, disable_gc
 from tinygrad.gradient import compute_gradient
 from tinygrad.mixin import OpMixin
@@ -25,7 +25,8 @@ def canonicalize_device(device:str|tuple|list|None) -> str|tuple[str, ...]:
 # *** all in scope Tensors are here. this gets relevant UOps ***
 
 all_tensors: dict[weakref.ref[Tensor], None] = {}
-def _apply_map_to_tensors(applied_map:dict[UOp, UOp], name:str, walk:bool=False) -> None:
+_pending_assigns: dict[UOp, list[UOp]] = {}  # buffer_uop -> [assign_uops in insertion order]
+def _apply_map_to_tensors(applied_map:dict[UOp, UOp], name:str) -> None:
   with cpu_profile(TracingKey(name), "TINY"):
     # get tensors in scope
     in_scope: dict[UOp, bool] = {}
@@ -34,7 +35,7 @@ def _apply_map_to_tensors(applied_map:dict[UOp, UOp], name:str, walk:bool=False)
 
     # get all Tensors and apply the map
     sink = UOp.sink(*[t.uop for t in scope_tensors])
-    new_sink = sink.substitute(applied_map, name=f"substitute {name}", walk=walk)
+    new_sink = sink.substitute(applied_map, name=f"substitute {name}")
 
     # set the relevant uop to the realized UOps
     for t,s,ns in zip(scope_tensors, sink.src, new_sink.src):
@@ -277,6 +278,23 @@ class Tensor(OpMixin):
   @disable_gc()
   def realize(self, *lst:Tensor, do_update_stats=True) -> Tensor:
     """Triggers the computation needed to create these Tensor(s)."""
+    # side-realize pending assigns for buffers referenced by these tensors
+    if _pending_assigns:
+      def _realize_pending(buf):
+        for assign_uop in _pending_assigns.pop(buf, []):
+          # recursively realize pending assigns that this assign's value depends on
+          for u in assign_uop.toposort():
+            if u.op is Ops.BUFFER and u in _pending_assigns: _realize_pending(u)
+          big_sink, becomes_map = transform_to_call(UOp.sink(assign_uop))
+          schedule, var_vals = complete_create_schedule_with_vars(big_sink)
+          _apply_map_to_tensors(becomes_map, name="Apply Pending Assign")
+          run_schedule(schedule, var_vals, do_update_stats=do_update_stats)
+          # update remaining pending assigns so they reference realized buffers instead of stale lazy graphs
+          if becomes_map:
+            for assigns in _pending_assigns.values():
+              for i in range(len(assigns)): assigns[i] = assigns[i].substitute(becomes_map)
+      for buf in {u for t in (self,)+lst for u in t.uop.toposort() if u.op is Ops.BUFFER}:
+        if buf in _pending_assigns: _realize_pending(buf)
     if len(to_realize:=[x for x in (self,)+lst if not x.uop.has_buffer_identity()]):
       run_schedule(*Tensor.schedule_with_vars(*to_realize), do_update_stats=do_update_stats)
     return self
@@ -305,13 +323,13 @@ class Tensor(OpMixin):
     if is_disk:
       self._buffer().copyin(x._data())
       return self
-    # NOTE: assign_uop is created before AFTER embedding (uses original self.uop),
-    # but AFTER must be embedded before _apply_uop (so subsequent assigns see it)
-    assign_uop = self.uop.assign(x.uop)
-    base = self.uop.base
-    if base.op in {Ops.BUFFER, Ops.AFTER} and not self.uop.has_buffer_identity():
-      _apply_map_to_tensors({base: base.after(assign_uop)}, name="Embed View Assign", walk=True)
-    return self.replace(self._apply_uop(lambda *_: assign_uop, x))
+    result = self._apply_uop(UOp.assign, x)
+    # track view assigns (not full-buffer or assign-chain) so they can be side-realized when the buffer is read
+    if (buf_uop:=self.uop.base).op is Ops.BUFFER and self.uop.op is not Ops.ASSIGN and not self.uop.has_buffer_identity():
+      # deduplicate: if the value is already a pending assign for this buffer (e.g. __iadd__ in __setitem__), remove it
+      if x.uop in _pending_assigns.get(buf_uop, []): _pending_assigns[buf_uop].remove(x.uop)
+      _pending_assigns.setdefault(buf_uop, []).append(result.uop)
+    return self.replace(result)
 
   def detach(self) -> Tensor:
     """
@@ -1218,6 +1236,26 @@ class Tensor(OpMixin):
     x_dims = [p for p in indices_parsed if not isinstance(p['index'], sint)]
     x = x.reshape(tuple(p['size'] for p in x_dims))
 
+    # basic setitem: construct result with view region replaced by v using arange masks
+    if v is not None and not any(isinstance(p['index'], Tensor) for p in indices_parsed):
+      # broadcast v to getitem shape, reshape to self.ndim (squeeze None dims, unsqueeze int dims — all are size 1)
+      vb = v.cast(self.dtype)._broadcast_to(x.shape)
+      vb = vb.reshape(tuple(1 if isinstance(p['index'], sint) else p['size'] for p in indices_parsed if p['index'] is not None))
+      # undo movement ops per-dim and build boolean mask
+      per_dim = []
+      for d, m in enumerate(mops):
+        (s, e), st = m['boundary'], abs(m['stride'])
+        if st != 1 and vb.shape[d] > 1:  # un-stride: interleave with zeros
+          vb = vb.unsqueeze(d+1)
+          vb = vb.pad_to(tuple(st if j == d+1 else None for j in range(vb.ndim)))
+          vb = vb.reshape(vb.shape[:d] + (vb.shape[d]*vb.shape[d+1],) + vb.shape[d+2:])
+          vb = vb.shrink_to(tuple(e-s if j == d else None for j in range(self.ndim)))
+        idx = Tensor.arange(self.shape[d], device=self.device).reshape([1]*d + [self.shape[d]] + [1]*(self.ndim - d - 1))
+        per_dim.append((idx >= s) & (idx < e) & (((e-1-idx) if m['stride'] < 0 else (idx-s)) % st == 0))
+      vb = vb.flip(tuple(d for d, m in enumerate(mops) if m['stride'] < 0))
+      vb = vb.pad(tuple((m['boundary'][0], self.shape[d] - m['boundary'][1]) for d, m in enumerate(mops)))
+      return (functools.reduce(lambda a, b: a & b, per_dim) if per_dim else Tensor(True, dtype=dtypes.bool, device=self.device)).where(vb, self)
+
     # tensor indexing
     if tops := [(d, p) for d, p in enumerate(x_dims) if isinstance(p['index'], Tensor)]:
       dims, tensors, masks = [d for d, _ in tops], cast(list[Tensor], [p['index'] for _, p in tops]), []
@@ -1228,7 +1266,7 @@ class Tensor(OpMixin):
       if v is None and len(dims) > 1 and consecutive and all_int(ishp := tuple(x.shape[d] for d in dims)):
         strides = tuple(prod(ishp[i+1:]) for i in range(len(dims)))
         try: linear_idx = functools.reduce(Tensor.add, (t._broadcast_to(big_shape) * s for t, s in zip(tensors, strides)))
-        except ValueError as err: raise IndexError(f"cannot broadcast indices: {err}") from err
+        except ValueError as e: raise IndexError(f"cannot broadcast indices: {e}") from e
         valid = functools.reduce(Tensor.__and__, ((t >= 0) & (t < s) for t, s in zip(tensors, ishp)))
         pre, post = x.shape[:dims[0]], x.shape[dims[-1]+1:]
         x = x.reshape(pre + (prod(ishp),) + post)[tuple([slice(None)] * len(pre)) + (valid.where(linear_idx, 0),)]
@@ -1239,7 +1277,7 @@ class Tensor(OpMixin):
       # create index masks
       for dim, tensor in zip(dims, tensors):
         try: i = tensor.reshape(tensor.shape + (1,)*(x.ndim - dims[0])).expand(pre_reduce_shape)
-        except ValueError as err: raise IndexError(f"cannot broadcast indices: {err}") from err
+        except ValueError as e: raise IndexError(f"cannot broadcast indices: {e}") from e
         masks.append(i._one_hot_along_dim(num_classes=x.shape[dim], dim=(dim - x.ndim)))
 
       # reduce masks to 1 mask
@@ -1248,36 +1286,21 @@ class Tensor(OpMixin):
       # inject 1's for the extra dims added in create masks
       reshape_arg = x.shape[:dims[0]] + (1,) * len(big_shape) + x.shape[dims[0]:]
       # sum reduce the extra dims introduced in create masks
-      x_pre = x  # save collapsed shape for advanced setitem
       x = (mask.where(x.reshape(reshape_arg), 0)).sum(sum_axis:=tuple(d + len(big_shape) for d in dims), dtype=x.dtype)
 
       # special permute case
       if (permuted := dims[0] != 0 and len(dims) != 1 and tuple(dims) != tuple(range(dims[0], dims[-1]+1))):
         mask, x = (y.permute(*range(dims[0], dims[0]+len(big_shape)), *range(0, dims[0]), *range(dims[0]+len(big_shape), y.ndim)) for y in (mask, x))
 
-      if v is None: return x  # advanced getitem
-      # advanced setitem: resolve tensor dims in collapsed space, then fall through to basic setitem path
-      vb = v.cast(self.dtype)._broadcast_to(_broadcast_shape(x.shape, v.shape))
-      for dim in sum_axis: vb = vb.unsqueeze(dim)  # add back reduced dims from sum
-      start = dims[0] if not permuted else 0
-      vb = _masked_setitem(x_pre, vb, mask, tuple(range(start, start + len(big_shape))))
-    elif v is None: return x  # basic getitem
-    # basic setitem: broadcast v, reshape to self.ndim (unsqueeze int dims, squeeze None dims)
-    else: vb = v.cast(self.dtype)._broadcast_to(x.shape)
-    vb = vb.reshape(tuple(1 if isinstance(p['index'], sint) else p['size'] for p in indices_parsed if p['index'] is not None))
-    per_dim = []
-    for d, m in enumerate(mops):
-      (s, e), st = m['boundary'], abs(m['stride'])
-      if st != 1 and vb.shape[d] > 1:  # un-stride: interleave with zeros
-        vb = vb.unsqueeze(d+1)
-        vb = vb.pad_to(tuple(st if j == d+1 else None for j in range(vb.ndim)))
-        vb = vb.reshape(vb.shape[:d] + (vb.shape[d]*vb.shape[d+1],) + vb.shape[d+2:])
-        vb = vb.shrink_to(tuple(e-s if j == d else None for j in range(self.ndim)))
-      idx = Tensor.arange(self.shape[d], device=self.device).reshape([1]*d + [self.shape[d]] + [1]*(self.ndim - d - 1))
-      per_dim.append((idx >= s) & (idx < e) & (((e-1-idx) if m['stride'] < 0 else (idx-s)) % st == 0))
-    vb = vb.flip(tuple(d for d, m in enumerate(mops) if m['stride'] < 0))
-    vb = vb.pad(tuple((m['boundary'][0], self.shape[d] - m['boundary'][1]) for d, m in enumerate(mops)))
-    return (functools.reduce(lambda a, b: a & b, per_dim) if per_dim else Tensor(True, dtype=dtypes.bool, device=self.device)).where(vb, self)
+      # for advanced setitem, returns whole tensor with indices replaced
+      if v is not None:
+        vb = v.cast(self.dtype)._broadcast_to(_broadcast_shape(x.shape, v.shape))
+        # add back reduced dims from sum
+        for dim in sum_axis: vb = vb.unsqueeze(dim)
+        # run _masked_setitem on tuple of axis that is to be reduced to match self.shape
+        x = _masked_setitem(self, vb, mask, tuple(range((start := dims[0] if not permuted else 0), start + len(big_shape))))
+
+    return x
 
   def __getitem__(self, indices) -> Tensor:
     """
@@ -1321,26 +1344,15 @@ class Tensor(OpMixin):
 
   def __setitem__(self, indices, v:Tensor|PyConst|list|tuple) -> None:
     if isinstance(v, Tensor) and v.dtype != self.dtype: raise RuntimeError(f"setitem dtype mismatch: {self.dtype=} != {v.dtype=}")
-    if self.requires_grad or (isinstance(v, Tensor) and v.requires_grad):
-      # for +=/-=, v's graph references self.uop through the view — exclude those from the stale-use check
-      v_uop, v_bw = (v.uop, v.uop.backward_slice) if isinstance(v, Tensor) else (None, {})
-      if any(self.uop in t.uop.backward_slice for tref in all_tensors
-             if (t:=tref()) is not None and t is not self and t.uop is not v_uop and t.uop not in v_bw):
-        raise RuntimeError("can't setitem on a tensor that already has other uses and requires grad")
-      if not isinstance(v, Tensor): v = Tensor(v, device=self.device, dtype=self.dtype)
-      if v.uop.op is Ops.ASSIGN: v = v._apply_uop(lambda x: x.src[1])
-      self.replace(self._getitem(indices, v))
-      return
+    if self.requires_grad or (isinstance(v, Tensor) and v.requires_grad): raise NotImplementedError("setitem with requires_grad is not supported")
     idx = [indices] if (isinstance(indices, list) and all_int(indices)) or not isinstance(indices, (tuple, list)) else list(indices)
     is_disk = isinstance(self.device, str) and self.device.startswith("DISK")
     if any(isinstance(i, (Tensor, list, tuple)) for i in idx): # advanced setitem
       if is_disk: raise RuntimeError("advanced setitem is not supported for DISK tensors")
       if not isinstance(v, Tensor): v = Tensor(v, device=self.device, dtype=self.dtype)
       self.assign(self._getitem(indices, v))
-    elif is_disk or self.uop.is_realized or self.uop.base.op is Ops.AFTER: # basic setitem, self is realized
-      view = self[indices]
-      if isinstance(v, Tensor) and v.uop.op is Ops.ASSIGN and v.uop in view.uop.base.src: return
-      view.assign(v)
+    elif is_disk or self.uop.is_realized: # basic setitem, self is realized. TODO: disk uop.base is a COPY and not realized
+      self[indices].assign(v)
     else: # basic setitem, self is not realized
       if not isinstance(v, Tensor): v = Tensor(v, device=self.device, dtype=self.dtype)
       # __iadd__/__isub__ on unrealized views creates a no-op ASSIGN; unwrap to get the computed value
@@ -1750,7 +1762,7 @@ class Tensor(OpMixin):
     print(t.all(axis=1, keepdim=True).numpy())
     ```
     """
-    return self.bool().min(axis, keepdim)
+    return self.logical_not().any(axis, keepdim).logical_not()
 
   def isclose(self, other:Tensor, rtol:float=1e-05, atol:float=1e-08, equal_nan=False) -> Tensor:
     """
@@ -2565,10 +2577,11 @@ class Tensor(OpMixin):
     return values._inverse(), indices
 
   @staticmethod
-  def _tri(r:sint, c:sint, diagonal=0, device=None, requires_grad:bool|None=None) -> Tensor:
+  def _tri(r:sint, c:sint, diagonal:int=0, device=None, requires_grad:bool|None=None) -> Tensor:
+    assert isinstance(r, int) and isinstance(c, int), f"does not support symbolic, getting {r=}, {c=}"
     return (Tensor.arange(r, device=device).unsqueeze(-1) + diagonal <= Tensor.arange(c, device=device)).requires_grad_(requires_grad)
 
-  def triu(self, diagonal:sint=0) -> Tensor:
+  def triu(self, diagonal:int=0) -> Tensor:
     """
     Returns the upper triangular part of the tensor, the other elements are set to 0.
 
@@ -2591,7 +2604,7 @@ class Tensor(OpMixin):
     """
     return Tensor._tri(self.shape[-2], self.shape[-1], diagonal=diagonal, device=self.device).where(self, self.zeros_like())
 
-  def tril(self, diagonal:sint=0) -> Tensor:
+  def tril(self, diagonal:int=0) -> Tensor:
     """
     Returns the lower triangular part of the tensor, the other elements are set to 0.
 
@@ -3265,7 +3278,7 @@ class Tensor(OpMixin):
     ```
     """
     if not dtypes.is_int(self.dtype): raise RuntimeError(f"expect integer dtype, getting {self.dtype=}")
-    if num_classes == -1: num_classes = int(self.max().item())+1
+    if num_classes == -1: num_classes = int((self.max()+1).item())
     return self[..., None]._one_hot_along_dim(num_classes).where(1, 0)
 
   def scaled_dot_product_attention(self, key:Tensor, value:Tensor, attn_mask:Tensor|None=None, dropout_p:float=0.0,
@@ -3283,6 +3296,9 @@ class Tensor(OpMixin):
     print(q.scaled_dot_product_attention(k, v).numpy())
     ```
     """
+    # NOTE: it also works when `key` and `value` have symbolic shape.
+    assert all_int(self.shape), f"does not support symbolic shape {self.shape}"
+
     if getenv("FLASH_ATTENTION"):
       from extra.thunder.tiny.fa import flash_attention
       return flash_attention(self, key, value, attn_mask=attn_mask, is_causal=is_causal)
@@ -3458,9 +3474,8 @@ class Tensor(OpMixin):
     #preprocess the matrix
     Q, R = (self.qr() if m >= n else self.transpose(-2, -1).qr())
     num, q_num = min(m, n), max(m, n)
-    # TODO: codegen infinite loop without contiguous
-    U = R.shrink(tuple([None] * len(b_shape) + [(0, num), (0, num)])).contiguous()
-    V = Tensor.eye(num, dtype=self.dtype).reshape((1,) * len(b_shape) + (num, num)).expand(b_shape + (num, num)).contiguous()
+    U = R.shrink(tuple([None] * len(b_shape) + [(0, num), (0, num)]))
+    V = Tensor.eye(num, dtype=self.dtype).reshape((1,) * len(b_shape) + (num, num)).expand(b_shape + (num, num))
     #prepare round robin pairing
     permute, inverse_permute = Tensor.arange(0, num, dtype=dtypes.int), Tensor.zeros(num, dtype=dtypes.int)
     permute[num//2:num] = permute[num//2:num].flip(0)
@@ -3598,7 +3613,7 @@ class Tensor(OpMixin):
     return cx.image_conv2d(cw, groups=groups, dtype=dtype).reshape(out_shape_t).transpose(self.ndim-1, self.ndim-2)
 
   def image_conv2d(self, weight:Tensor, bias:Tensor|None=None, groups=1, stride=1, dilation=1, padding=0, dtype=None) -> Tensor:
-    base_image_type, dtsz = (dtypes.imageh, 2) if FLOAT16 else (dtypes.imagef, 4)
+    base_image_type, dtsz = (dtypes.imageh, 2) if (FLOAT16:=getenv("FLOAT16", 0)) else (dtypes.imagef, 4)
 
     (bs,_,iy,ix), (cout,cin,H,W) = self.shape, weight.shape
     x, w = self, weight.reshape(groups, (rcout := cout//groups), cin, H, W)

tinygrad/uop/ops.py

@@ -26,7 +26,7 @@ axis_colors = {AxisType.GLOBAL: "blue", AxisType.THREAD: "BLUE", AxisType.LOCAL:
 axis_to_pos = {AxisType.LOOP: -1, AxisType.THREAD: 0, AxisType.GLOBAL: 0, AxisType.WARP: 1, AxisType.LOCAL: 2, AxisType.UPCAST: 3,
                AxisType.GROUP_REDUCE: 2, AxisType.REDUCE: 4, AxisType.UNROLL: 5}
 
-range_start = {Ops.BUFFERIZE: 1, Ops.REDUCE: 1, Ops.STORE: 2, Ops.WMMA: 3, Ops.END: 1, Ops.CALL: 1, Ops.COPY: 2, Ops.ALLREDUCE: 2, Ops.BUFFER_VIEW: 1}
+range_start = {Ops.BUFFERIZE: 1, Ops.REDUCE: 1, Ops.STORE: 2, Ops.WMMA: 3, Ops.END: 1, Ops.CALL: 1, Ops.COPY: 2, Ops.BUFFER_VIEW: 1}
 
 # https://en.wikipedia.org/wiki/Identity_element
 def identity_element(op:Ops, dt:DType) -> PyConst: return dtypes.as_const({Ops.ADD:0, Ops.MUL:1, Ops.MAX:dtypes.min(dt)}[op], dt)
@@ -163,7 +163,7 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
     # Check self first, then iterate backward_slice (avoids creating intermediate dict)
     return self.op in ops or any(x.op in ops for x in self.backward_slice)
 
-  def toposort(self, gate:Callable|None=None, enter_calls=True) -> dict[UOp, None]:
+  def toposort(self, gate:Callable|None=None) -> dict[UOp, None]:
     cache: dict[UOp, None] = {}
     stack: list[tuple[UOp, bool]] = [(self, False)] # each stack entry is (node, visited_flag)
     while stack:
@@ -172,8 +172,7 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
       if not visited:
         if gate is None or gate(node):
           stack.append((node, True))  # push node back on stack to process after its srcs
-          for s in reversed(node.src if enter_calls or node.op is not Ops.CALL else node.src[1:]):
-            stack.append((s, False)) # push srcs on the stack
+          for s in reversed(node.src): stack.append((s, False)) # push srcs on the stack
       else: cache[node] = None # second time i'm seeing this node, add it to returned toposort
     return cache
 
@@ -254,9 +253,6 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
       case Ops.RESHAPE:
         if self.src[0]._shape is None: return self.marg
 
-      # MULTI marker (axis info in PARAM sources) has no shape
-      case Ops.MULTI if len(self.src) == 0: return None
-
     # movement ops change the shape
     # NOTE: ssimplify is required because the shape needs to be canonical for broadcasting and same shape checking
     if self.op in GroupOp.Movement.union({Ops.MULTI, Ops.REDUCE_AXIS, Ops.WMMA}):
@@ -518,11 +514,6 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
     # COPY removes axis. TODO: add more tests for this, and consider MSELECT/MSTACK
     if self.op is Ops.COPY: return None
     if self.op is Ops.MULTI: return self.arg
-    # PARAM: axis is stored as a MULTI source
-    if self.op is Ops.PARAM:
-      for s in self.src:
-        if s.op is Ops.MULTI: return s.arg
-      return None
     # NOTE: they all have to share an axis, we always choose [-1]
     if self.op in GroupOp.ALU: return axes[-1] if (axes := dedup([x.axis for x in self.src if x.axis is not None])) else None
     if len(self.src) == 0: return None
@@ -657,45 +648,34 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
     while len(s.src) and s.op not in {Ops.BUFFER, Ops.PARAM, Ops.BUFFERIZE, Ops.MSTACK}: s = s.src[0]
     return s
 
-  def contiguous_view_offset(self) -> int|None:
-    """If movement ops on a BUFFER collapse to a contiguous range, return `offset` in elements. Otherwise None."""
-    from tinygrad.schedule.rangeify import pm_mops
-    from tinygrad.uop.symbolic import symbolic
-    out = graph_rewrite(self._mop(Ops.RESHAPE, (self.size,)).index(UOp.range(self.size, 0)), pm_mops+symbolic, name="contiguous_view_offset")
-    if out.op is not Ops.INDEX: return None
-    if out.src[1].op is Ops.CONST and self.size == 1:
-      if not isinstance(out.src[1].arg, int): return None  # masked/padded regions produce InvalidType
-      return out.src[1].arg
-    if out.src[1].op is Ops.RANGE: return 0
-    if out.src[1].op is Ops.ADD and out.src[1].src[0].op is Ops.RANGE and out.src[1].src[1].op is Ops.CONST:
-      if not isinstance(out.src[1].src[1].arg, int): return None  # masked/padded regions produce InvalidType
-      return out.src[1].src[1].arg
-    return None
-
   def has_buffer_identity(self):
     """Check if this UOp has a concrete buffer identity in the graph (RESHAPE/MULTI -> BUFFER chain)."""
     if self.op in {Ops.RESHAPE, Ops.MULTI}: return self.src[0].has_buffer_identity()
-    return self.op in {Ops.BUFFER, Ops.BUFFER_VIEW, Ops.PARAM}
+    return self.op in {Ops.BUFFER, Ops.PARAM}
 
   @property
   def buffer(self) -> Buffer|MultiBuffer:
     from tinygrad.device import Buffer, MultiBuffer
-    if self.op in {Ops.CONTIGUOUS, Ops.RESHAPE, Ops.DETACH, Ops.AFTER}: return self.src[0].buffer
+    if self.op in {Ops.CONTIGUOUS, Ops.RESHAPE}: return self.src[0].buffer
     # this buffer can process disk tensors and simple movement ops
     if self is not self.base:
-      offset = self.contiguous_view_offset()
-      if offset is None: raise RuntimeError(f"cannot collapse movement ops on {self.base.op} to a contiguous view")
-      buf = self.base.buffer
+      from tinygrad.schedule.rangeify import pm_mops
+      from tinygrad.uop.symbolic import symbolic
+      out = graph_rewrite(self.flatten().index(UOp.range(self.size, 0)), pm_mops+symbolic)
+      buf = out.src[0].buffer
       assert isinstance(buf, Buffer), "must be a Buffer for movement ops"
-      return buf.view(self.size, self.dtype, offset*self.dtype.itemsize)
+      assert out.op is Ops.INDEX, "couldn't collapse to a single INDEX"
+      if out.src[1].op is Ops.CONST:
+        return buf.view(1, out.dtype, out.src[1].arg*out.dtype.itemsize)
+      if out.src[1].op is Ops.RANGE:
+        return buf.view(self.size, out.dtype, 0)
+      if out.src[1].op is Ops.ADD and out.src[1].src[0].op is Ops.RANGE and out.src[1].src[1].op is Ops.CONST:
+        return buf.view(self.size, out.dtype, out.src[1].src[1].arg*out.dtype.itemsize)
+      raise RuntimeError(f"cannot collapse INDEX {out.pyrender()} to a single size/offset")
     if self.op is Ops.BITCAST:
       buf = self.src[0].buffer
       assert isinstance(buf, Buffer), "must be a Buffer for BITCAST"
       return buf.view(self.size, self.dtype, 0)
-    if self.op is Ops.BUFFER_VIEW:
-      buf = self.src[0].buffer
-      assert isinstance(buf, Buffer), "must be a Buffer for BUFFER_VIEW"
-      return buf.view(self.size, self.dtype, self.arg[1] * self.dtype.itemsize)
     if self.op is Ops.MSELECT:
       ret = self.src[0].buffer
       assert isinstance(ret, MultiBuffer)
@@ -887,14 +867,12 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
   def param_like(self, slot:int):
     if self.op is Ops.BIND:
       return UOp.param(slot, self.dtype, self._shape, self._device, self._min_max, self.src[0].arg[0])
-    p = UOp.param(slot, self.dtype, self._shape, self._device)
-    if self.axis is not None: p = p.replace(src=p.src + (UOp(Ops.MULTI, arg=self.axis),))
-    return p
+    return UOp.param(slot, self.dtype, self._shape, self._device)
 
-  def call(self, *srcs:UOp, grad_fxn:Callable|None=None, metadata:tuple[Metadata, ...]=(), name:str|None=None, precompile:bool=False) -> UOp:
+  def call(self, *srcs:UOp, grad_fxn:Callable|None=None, metadata:tuple[Metadata, ...]=(), name:str|None=None) -> UOp:
     # TODO: reenable this after ENCDEC is fixed
     #assert len(self.ranges) == 0, f"ranges {self.ranges} are leaking out of the call in {self.pyrender()}"
-    return UOp(Ops.CALL, self.dtype, (self,)+srcs, CallInfo(grad_fxn, metadata, name, precompile))
+    return UOp(Ops.CALL, self.dtype, (self,)+srcs, CallInfo(grad_fxn, metadata, name))
   def custom_kernel(*srcs:UOp, fxn:Callable, grad_fxn:Callable|None=None) -> list[UOp]:
     contig_srcs = tuple(x.contiguous() if x.op is not Ops.AFTER else x for x in srcs)
     placeholders = [UOp.placeholder_like(s, slot=i) for i,s in enumerate(contig_srcs)]
@@ -917,15 +895,15 @@ class CallInfo:
   grad_fxn: Callable|None = None
   metadata: tuple[Metadata, ...] = ()
   name: str|None = None
-  precompile: bool = False
   # grad_fxn can't be pickled, but metadata can
-  def __reduce__(self): return (CallInfo, (None, self.metadata, self.name, self.precompile))
-  def __repr__(self): return f"CallInfo({id(self.grad_fxn) if self.grad_fxn else None}, {self.metadata}, {repr(self.name)}, {self.precompile})"
+  def __reduce__(self): return (CallInfo, (None, self.metadata, self.name))
+  def __repr__(self): return f"CallInfo({id(self.grad_fxn) if self.grad_fxn else None}, {self.metadata}, {repr(self.name)})"
 
 def should_resolve_call(c:UOp) -> bool:
   # 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 False
-  if c.src[0].op in {Ops.PROGRAM, Ops.LINEAR, Ops.COPY, Ops.ALLREDUCE}: return False
+  if c.src[0].op is Ops.PROGRAM: return False
+  if c.src[0].op is Ops.COPY: return False
   return True
 
 # ******** ops in python ********
@@ -1272,13 +1250,12 @@ if TRACK_MATCH_STATS or PROFILE:
 SENTINEL: Final[UOp] = cast(UOp, object())
 class BottomUpGate(Exception): pass
 class RewriteContext:
-  def __init__(self, pm, bpm, ctx=None, enter_calls=False):
+  def __init__(self, pm, bpm, ctx=None):
     self.pm: PatternMatcher|None = pm
     self.bpm: PatternMatcher|None = bpm
     self.bpm_cache: dict[UOp, UOp|None] = {}
     self.ctx = ctx
     self.replace: dict[UOp, UOp] = {}
-    self.enter_calls = enter_calls
 
   # no cache needed: pm_rewrite is called at most once per UOp due to the replace dict check in unified_rewrite
   def pm_rewrite(self, x:UOp) -> UOp|None: return unwrap(self.pm).rewrite(x, self.ctx)
@@ -1301,7 +1278,7 @@ class RewriteContext:
           continue
         # no rewrite, process children then come back to rebuild
         stack.append((n, True))
-        if not self.enter_calls and n.op is Ops.CALL: self.replace[n.src[0]] = n.src[0]
+        if n.op is Ops.CALL: self.replace[n.src[0]] = n.src[0]
         for x in reversed(n.src):
           if x not in self.replace: stack.append((x, False))
       else:
@@ -1341,7 +1318,7 @@ class RewriteContext:
         # NOTE: CALL is handled as a special case.
         # The function that is called is not included in the graph_rewrite.
         # If you want to graph_rewrite a call, you can
-        if not self.enter_calls and new_n.op is Ops.CALL: self.replace[new_n.src[0]] = new_n.src[0]
+        if new_n.op is Ops.CALL: self.replace[new_n.src[0]] = new_n.src[0]
         for x in reversed(new_n.src):
           if x in on_stack: continue
           stack.append((x, 0, x))
@@ -1380,8 +1357,8 @@ class RewriteContext:
     return self.replace[root]
 
 @profile_matches
-def graph_rewrite(sink:UOp, pm:PatternMatcher, ctx=None, bottom_up=False, name=None, bpm=None, walk=False, enter_calls=False) -> UOp:
-  rewrite_ctx = RewriteContext(pm if not bottom_up else None, pm if bottom_up else bpm, ctx, enter_calls)
+def graph_rewrite(sink:UOp, pm:PatternMatcher, ctx=None, bottom_up=False, name=None, bpm=None, walk=False) -> UOp:
+  rewrite_ctx = RewriteContext(pm if not bottom_up else None, pm if bottom_up else bpm, ctx)
   return rewrite_ctx.walk_rewrite(sink) if walk else rewrite_ctx.unified_rewrite(sink)
 
 def sint_to_uop(x:sint, dtype=dtypes.index) -> UOp: return UOp.const(dtype, x) if isinstance(x, int) else x.cast(dtype)
@@ -1415,10 +1392,7 @@ def _index_to_concrete_int(u:UOp) -> UOp: return graph_rewrite(u.sink(), pm_lowe
 _substitute = PatternMatcher([(UPat(tuple(Ops), name="x"), lambda ctx,x: ctx.get(x,None))])
 _remove_all_tags = PatternMatcher([(UPat(GroupOp.All, name="x"), lambda x: x.replace(tag=None) if x.tag is not None else None)])
 
-def gate_kernel_sink(x:UOp) -> bool:
-  if x.op is Ops.LINEAR: return False
-  if x.op is Ops.SINK and isinstance(x.arg, KernelInfo): return False
-  return True
+def gate_kernel_sink(x:UOp) -> bool: return not (x.op is Ops.SINK and isinstance(x.arg, KernelInfo))
 
 def do_unbind(ctx:dict[Variable, int], x:UOp):
   v,i = x.unbind()
@@ -1448,7 +1422,6 @@ def bitcast(x, in_dtype:DType, out_dtype:DType):
 
 renderer = PatternMatcher([
   (UPat((Ops.DEFINE_VAR,), name="x"), lambda x: x.expr),
-  (UPat(Ops.PARAM, src=(UPat(), UPat(), UPat(), UPat(), UPat(Ops.NOOP, name="x"))), lambda x: x.arg),
   (UPat((Ops.SPECIAL), name="x"), lambda x: x.arg),
   (UPat(Ops.RANGE, name="x"), lambda x: f"r{range_str(x)}"),
   (UPat((Ops.CONST, Ops.VCONST), name="x"), lambda x: str(x.arg)),
@@ -1510,8 +1483,7 @@ pm_pyrender_extra = PatternMatcher([
   (UPat(Ops.INDEX, src=(UPat(), UPat()), allow_any_len=True, name="x"), lambda ctx,x:
    f"{ctx[x.src[0]]}.index({ctx[x.src[1]]}, "+(f"{ctx[x.src[2]]}, " if len(x.src) > 2 else "")+
     (f"dtype={x.dtype})" if x.src[0].dtype != x.dtype else "ptr=True)") if x.src[0].dtype.base != x.dtype else None),
-  # TODO: movement ops simplify stuff, this can break SPEC=2
-  #(UPat(GroupOp.Movement, name="x"), lambda ctx,x: f"{ctx[x.src[0]]}.{x.op.name.lower()}({render_marg(ctx,x)})"),
+  (UPat(GroupOp.Movement, name="x"), lambda ctx,x: f"{ctx[x.src[0]]}.{x.op.name.lower()}({render_marg(ctx,x)})"),
   # NOTE: CMPNE doesn't work cause there's no __rne__
   # NOTE: only match CONSTs without UNIQUE (len(src)==1), unique_const needs explicit rendering
   (UPat(set(syms.keys())-{Ops.SUB, Ops.CMPNE}, src=(UPat(Ops.CONST, src=(UPat(Ops.DEVICE),), name="y"), UPat(name="z")), name="x"),
@@ -1537,7 +1509,7 @@ def pyrender(ast:UOp) -> str:
   cmap = consumer_map_from_toposort(lst)
   not_rendered = {Ops.CONST, Ops.VCONST, Ops.DEVICE}
   always_rendered = {Ops.PARAM, Ops.LOAD, Ops.SPECIAL, Ops.RANGE, Ops.CONTIGUOUS, Ops.VECTORIZE,
-                     Ops.BUFFER, Ops.COPY, Ops.ALLREDUCE, Ops.CALL, Ops.WHERE, Ops.END, Ops.ASSIGN}
+                     Ops.BUFFER, Ops.COPY, Ops.CALL, Ops.WHERE, Ops.END, Ops.ASSIGN}
 
   to_render: set[UOp] = {ast}
   for u in lst:

tinygrad/uop/spec.py

@@ -87,9 +87,6 @@ _tensor_spec = PatternMatcher([
   (UPat(Ops.BUFFER, src=(UPat((Ops.LUNIQUE, Ops.UNIQUE)), UPat(Ops.DEVICE)), name="buf"),
    lambda buf: isinstance(buf.arg, int) and isinstance(buf.dtype, (DType, ImageDType))),
 
-  # BUFFER_VIEW on BUFFER is allowed if BUFFER is
-  (UPat(Ops.BUFFER_VIEW, src=(UPat(Ops.BUFFER),)), lambda: True),
-
   # KERNEL can attach to an AFTER to describe the compute required to realize a BUFFER
   (UPat(Ops.CALL, src=UPat((Ops.BUFFER, Ops.AFTER, Ops.MSELECT, Ops.MSTACK, Ops.BIND))), lambda: True),
 
@@ -209,11 +206,9 @@ kernel_spec = PatternMatcher([
   # reduce must be on ranges
   (UPat(Ops.REDUCE, src=(UPat(),), allow_any_len=True, name="x"), lambda x: all(y.dtype in (dtypes.index, dtypes.int) for y in x.src[1:])),
 
-  # COPY/ALLREDUCE/BUFFER_VIEW can have ranges appended
+  # COPY/BUFFER_VIEW can have ranges appended
   (UPat(Ops.COPY, name="x", src=(UPat.var("s"), UPat(Ops.DEVICE)), allow_any_len=True, arg=None),
    lambda x,s: x.dtype == s.dtype and all(u.op is Ops.RANGE for u in x.src[2:])),
-  (UPat(Ops.ALLREDUCE, name="x", src=(UPat.var("s"), UPat(Ops.DEVICE)), allow_any_len=True),
-   lambda x,s: x.dtype == s.dtype and isinstance(x.arg, Ops) and all(u.op is Ops.RANGE for u in x.src[2:])),
   (UPat(Ops.BUFFER_VIEW, src=(UPat((Ops.INDEX, Ops.LOAD)),), allow_any_len=True, name="x"),
    lambda x: all(u.op is Ops.RANGE for u in x.src[1:])),
 ])+movement_ops+shared_codegen_spec+shared_spec

tinygrad/uop/symbolic.py

@@ -427,6 +427,8 @@ sym = symbolic+pm_simplify_valid+PatternMatcher([
   (UPat((Ops.SINK, Ops.GROUP), name="root"),
     lambda root: UOp(root.op, root.dtype, tuple(flatten(x.src if x.op in REMOVE_FROM_SINK_LIKE else (x,) for x in root.src)), root.arg)
       if any(x.op in REMOVE_FROM_SINK_LIKE for x in root.src) else None),
+  # remove END with empty NOOP
+  (UPat(Ops.END, src=(UPat(Ops.NOOP, src=(), name="noop"),), allow_any_len=True), lambda noop:noop),
   # ** combine terms (opinionated) **
   (-1 * (UPat.var("x") + UPat.var("y")), lambda x,y: (-x)+(-y)),  # -(x+y) -> -x + -y
   # (x+y)*c -> x*c+y*c. only for int, float has inf*0=nan issue

tinygrad/viz/index.html

@@ -143,31 +143,6 @@
   g.label rect.bg.highlight {
     fill: #5f0059;
   }
-  #insts .line {
-    display: flex;
-    flex-direction: column;
-    cursor: pointer;
-    margin-bottom: 8px;
-  }
-  #insts .left {
-    width: fit-content;
-    display: flex;
-    gap: 4px;
-  }
-  #insts .left.highlight {
-    background-color: rgba(0, 199, 47, 0.2);
-  }
-  #insts .n {
-    color: #787fa1;
-    min-width: 5ch;
-  }
-  #insts .wave {
-    color: #7aa2f7;
-    min-width: 2ch;
-  }
-  #insts .pc {
-    color: #73daca;
-  }
   g.node.highlight rect.node, .edgePath.highlight, g.port circle {
     stroke: #89C9A2;
   }

tinygrad/viz/js/index.js

@@ -62,7 +62,7 @@ const drawGraph = (data) => {
   const callCount = g.graph().callCount;
   const nodes = d3.select("#nodes").selectAll("g").data(g.nodes().map(id => g.node(id)), d => d).join("g").attr("class", d => d.className ?? "node")
     .attr("transform", d => `translate(${d.x},${d.y})`).on("click", (e,d) => {
-      if (d.callNode) {
+      if (d.label.startsWith("CALL")) {
         if (state.callSrcMask.has(d.id)) state.callSrcMask.delete(d.id); else state.callSrcMask.add(d.id);
         if (state.callSrcMask.size >= callCount) { showCallSrc.toggle.checked = !showCallSrc.toggle.checked; state.callSrcMask.clear(); }
         return setState({});
@@ -110,7 +110,7 @@ const drawGraph = (data) => {
   });
   addTags(nodes.selectAll("g.tag").data(d => d.tag != null ? [d] : []).join("g").attr("class", "tag")
     .attr("transform", d => `translate(${-d.width/2+8}, ${-d.height/2+8})`).datum(e => e.tag));
-  addTags(nodes.selectAll("g.type").data(d => d.callNode ? [d] : []).join("g")
+  addTags(nodes.selectAll("g.type").data(d => d.label.startsWith("CALL\n") ? [d] : []).join("g")
     .attr("class", d => `tag ${d.collapsed ? 'collapsed' : 'expanded'}`)
     .attr("transform", d => `translate(${-d.width/2}, ${0})`).datum(d => d.collapsed ? "+" : "−"));
   // draw edges
@@ -234,7 +234,7 @@ const drawLine = (ctx, x, y, opts) => {
 }
 
 function tabulate(rows) {
-  const root = d3.create("div").style("display", "grid").style("grid-template-columns", `${Math.max(...rows.map(x => x[0].length), 0)}ch 1fr`).style("gap", "0.2em").style("white-space", "nowrap");
+  const root = d3.create("div").style("display", "grid").style("grid-template-columns", `${Math.max(...rows.map(x => x[0].length), 0)}ch 1fr`).style("gap", "0.2em");
   for (const [k,v] of rows) { root.append("div").text(k); root.append("div").node().append(v); }
   return root;
 }
@@ -253,19 +253,10 @@ const Modes = {0:'read', 1:'write', 2:'write+read'};
 function setFocus(key) {
   if (key !== focusedShape) {
     saveToHistory({ shape:focusedShape });
-    // adjust zoom if the entire shape is off screen
-    const { eventType, e } = selectShape(key);
-    if (e != null) {
-      const [x0, x1] = eventType === EventTypes.EXEC ? [e.x, e.x+e.width] : [e.x[0], e.x.at(-1)];
-      const xscale = d3.scaleLinear().domain([data.first, data.dur]).range([0, document.getElementById("timeline").clientWidth]);
-      const [st, et] = xscale.range().map(zoomLevel.invertX, zoomLevel).map(xscale.invert, xscale);
-      if (x1 < st || x0 > et) zoomLevel = d3.zoomIdentity.translate(-xscale((x0+x1)/2-(et-st)/2)*zoomLevel.k, 0).scale(zoomLevel.k);
-    }
     focusedShape = key; d3.select("#timeline").call(canvasZoom.transform, zoomLevel);
   }
   const { eventType, e } = selectShape(key);
-  if (metadata.querySelector(".info") == null) d3.select(metadata).html("").append("div").classed("info", true);
-  const html = d3.select(".info").html("");
+  const html = d3.create("div").classed("info", true);
   if (eventType === EventTypes.EXEC) {
     const [n, _, ...rest] = e.arg.tooltipText.split("\n");
     html.append(() => tabulate([["Name", d3.create("p").html(n).node()], ["Duration", formatTime(e.width)], ["Start Time", formatTime(e.x)]]).node());
@@ -299,24 +290,7 @@ function setFocus(key) {
       if (shape != null) p.style("cursor", "pointer").on("click", () => setFocus(shape));
     }
   }
-  // instructions list renderer
-  let instList = document.getElementById("insts");
-  if (data.pcToShape.size > 0 && instList == null) {
-    let contents = "", i = 0;
-    for (const [k, v] of data.pcToShape) {
-      contents += `<div class="line" data-k="${k}"><span class="left" id="inst-${k}"><span class="n">${i++}</span><span class="wave">${v.wave}</span>
-        <span class="pc">${"0x"+v.pc.toString(16).padStart(12, "0")}</span></span><span class="label">${data.pcMap[v.pc]}</span></div>`;
-    }
-    instList = d3.create("pre").append("code").classed("hljs", true).style("margin-top", "20px").attr("id", "insts").html(contents)
-      .on("click", e => { const line = e.target.closest(".line"); line && setFocus(line.dataset.k); }).node();
-    metadata.insertBefore(instList.parentElement, html.node());
-  }
-  d3.select(instList).selectAll("span").classed("highlight", false);
-  const instLine = document.getElementById(`inst-${key}`); instLine?.classList.add("highlight");
-  if (instLine != null && instList != null) {
-    const r = rect(instLine), c = rect(instList);
-    if (Math.max(c.top-r.bottom, r.top-c.bottom)>=-30) instLine.scrollIntoView({ block:"center" });
-  }
+  return metadata.replaceChildren(html.node());
 }
 
 const EventTypes = { EXEC:0, BUF:1 };
@@ -325,7 +299,7 @@ async function renderProfiler(path, unit, opts) {
   displaySelection("#profiler");
   // support non realtime x axis units
   formatTime = unit === "realtime" ? formatMicroseconds : formatCycles;
-  if (data?.path !== path) { data = {tracks:new Map(), axes:{}, path, first:null, pcToShape:new Map()}; focusedDevice = null; focusedShape = null; }
+  if (data?.path !== path) { data = {tracks:new Map(), axes:{}, path, first:null}; focusedDevice = null; focusedShape = null; }
   setFocus(focusedShape);
   // layout once!
   if (data.tracks.size !== 0) return updateProgress(Status.COMPLETE);
@@ -338,9 +312,9 @@ async function renderProfiler(path, unit, opts) {
   const u64 = () => { const ret = new Number(view.getBigUint64(offset, true)); offset += 8; return ret; }
   const f32 = () => { const ret = view.getFloat32(offset, true); offset += 4; return ret; }
   const optional = (i) => i === 0 ? null : i-1;
-  const dur = u32(), tracePeak = u64(), indexLen = u32(), layoutsLen = u32(); data.dur = dur;
+  const dur = u32(), tracePeak = u64(), indexLen = u32(), layoutsLen = u32();
   const textDecoder = new TextDecoder("utf-8");
-  const { strings, dtypeSize, markers, ...extData } = JSON.parse(textDecoder.decode(new Uint8Array(buf, offset, indexLen))); offset += indexLen;
+  const { strings, dtypeSize, markers }  = JSON.parse(textDecoder.decode(new Uint8Array(buf, offset, indexLen))); offset += indexLen;
   // place devices on the y axis and set vertical positions
   const [tickSize, padding, baseOffset] = [10, 8, markers.length ? 14 : 0];
   const deviceList = profiler.append("div").attr("id", "device-list").style("padding-top", tickSize+padding+baseOffset+"px");
@@ -359,8 +333,7 @@ async function renderProfiler(path, unit, opts) {
     const k = textDecoder.decode(new Uint8Array(buf, offset, nameLen)); offset += nameLen;
     const div = deviceList.append("div").attr("id", k).text(k).style("padding", padding+"px").style("width", opts.width);
     const { y:baseY, height:baseHeight } = rect(div.node());
-    const [dname, dnum] = k.split(":", 2);
-    const colors = colorScheme[dname] ?? colorScheme.DEFAULT;
+    const colors = colorScheme[k.split(":")[0]] ?? colorScheme.DEFAULT;
     const offsetY = baseY-canvasTop+padding/2;
     const shapes = [], visible = [];
     const eventType = u8(), eventsLen = u32();
@@ -402,10 +375,8 @@ async function renderProfiler(path, unit, opts) {
         if (shapeRef != null) { ref = {ctx:e.ref, step:0}; shapeRef = ref; }
         else if (ref != null) {
           const start = ref.step>0 ? ref.step+1 : 0;
-          const steps = ctxs[ref.ctx+1].steps;
-          for (let si=start; si<steps.length; si++) {
-            if (steps[si].name == e.name) { ref.step = si; shapeRef = ref; break; }
-          }
+          const stepIdx = ctxs[ref.ctx+1].steps.findIndex((s, i) => i >= start && s.name == e.name);
+          if (stepIdx !== -1) { ref.step = stepIdx; shapeRef = ref; }
         } else {
           const steps = ctxs[state.currentCtx].steps;
           for (let i=state.currentStep+1; i<steps.length; i++) {
@@ -417,9 +388,8 @@ async function renderProfiler(path, unit, opts) {
         // tiny device events go straight to the rewrite rule
         const key = k.startsWith("TINY") ? null : `${k}-${j}`;
         const labelHTML = label.map(l=>`<span style="color:${l.color}">${l.st}</span>`).join("");
-        let info = e.info != null ? "\n"+e.info : "";
-        if (info.startsWith("\nPC:")) data.pcToShape.set(key, {wave:dnum, pc:parseInt(e.info.split(":")[1]), st:e.st}); info = "";
-        const arg = { tooltipText:labelHTML+" N:"+shapes.length+"\n"+formatTime(e.dur)+info, bufs:[], key, ctx:shapeRef?.ctx, step:shapeRef?.step };
+        const arg = { tooltipText:labelHTML+" N:"+shapes.length+"\n"+formatTime(e.dur)+(e.info != null ? "\n"+e.info : ""), bufs:[], key,
+                      ctx:shapeRef?.ctx, step:shapeRef?.step };
         if (e.key != null) shapeMap.set(e.key, key);
         // offset y by depth
         shapes.push({x:e.st, y:levelHeight*depth, width:e.dur, height:levelHeight, arg, label:opts.hideLabels ? null : label, fillColor });
@@ -428,74 +398,65 @@ async function renderProfiler(path, unit, opts) {
       div.style("height", levelHeight*levels.length+padding+"px").style("pointerEvents", "none");
     } else {
       const peak = u64();
+      let x = 0, y = 0;
+      const buf_shapes = new Map(), temp = new Map();
       const timestamps = [], valueMap = new Map();
-      // start by unpacking the raw events
-      const memEvents = [];
-      let x = 0, y = 0, shapeIdx = 0;
-      const allocs = new Map();
       for (let j=0; j<eventsLen; j++) {
         const alloc = u8(), ts = u32(), key = u32();
         if (alloc) {
           const dtype = strings[u32()], sz = u64(), nbytes = dtypeSize[dtype]*sz;
-          allocs.set(key, {nbytes, shapeKey:`${k}-${shapeIdx++}`});
-          memEvents.push({alloc, key, dtype, sz, nbytes});
+          const shape = {x:[x], y:[y], dtype, sz, nbytes, key};
+          buf_shapes.set(key, shape); temp.set(key, shape);
           timestamps.push(ts);
           x += 1; y += nbytes; valueMap.set(ts, y);
         } else {
-          const users = Array.from({ length: u32() }, () => ({shape:shapeMap.get(u32()), repr:strings[u32()], num:u32(), mode:u8()}));
-          const {nbytes, shapeKey} = allocs.get(key); allocs.delete(key);
-          users?.forEach((u) => selectShape(u.shape).e?.arg.bufs.push({ key:shapeKey, nbytes, num:u.num, mode:u.mode, k }));
-          memEvents.push({alloc, key, users, nbytes});
+          const free = buf_shapes.get(key);
+          free.users = Array.from({ length: u32() }, () => ({shape:shapeMap.get(u32()), repr:strings[u32()], num:u32(), mode:u8()}));
           timestamps.push(ts); valueMap.set(ts, y);
-          x += 1; y -= nbytes;
+          x += 1; y -= free.nbytes;
+          free.x.push(x);
+          free.y.push(free.y.at(-1));
+          temp.delete(key);
+          for (const [k, v] of temp) {
+            if (k <= key) continue;
+            v.x.push(x, x);
+            v.y.push(v.y.at(-1), v.y.at(-1)-free.nbytes);
+          }
         }
       }
       timestamps.push(dur);
       const height = heightScale(peak);
       const yscale = d3.scaleLinear().domain([0, peak]).range([height, 0]);
+      for (const [num, {dtype, sz, nbytes, y, x:steps, users}] of buf_shapes) {
+        const x = steps.map(s => timestamps[s]);
+        const dur = x.at(-1)-x[0];
+        const arg = { tooltipText:`${dtype}\n${formatUnit(sz)}\n${formatUnit(nbytes, 'B')}\n${formatTime(dur)}`, users, key:`${k}-${shapes.length}` };
+        shapes.push({ x, y0:y.map(yscale), y1:y.map(y0 => yscale(y0+nbytes)), arg, fillColor:cycleColors(colorScheme.BUFFER, shapes.length) });
+        users?.forEach((u) => selectShape(u.shape).e?.arg.bufs.push({ key:arg.key, nbytes, num:u.num, mode:u.mode, k }));
+      }
       // generic polygon merger
       const base0 = yscale(0);
-      const sum = {x:[], y0:[], y1:[], fillColor:"#2B1B72"};
-      for (let i=0; i<timestamps.length-1; i++) {
-        const yv = yscale(valueMap.get(timestamps[i]));
-        sum.x.push(timestamps[i], timestamps[i+1]); sum.y1.push(yv, yv); sum.y0.push(base0, base0);
-      }
-      // build individual buffer shapes when user clicks to expand, this detailed layout is n²
-      let bufShapes = null;
-      const buildBufShapes = () => {
-        if (bufShapes != null) return bufShapes;
-        bufShapes = [];
-        const buf_shapes = new Map(), temp = new Map();
-        let x = 0, y = 0;
-        for (const e of memEvents) {
-          if (e.alloc) {
-            const shape = {x:[x], y:[y], dtype:e.dtype, sz:e.sz, nbytes:e.nbytes, key:e.key};
-            buf_shapes.set(e.key, shape); temp.set(e.key, shape);
-            x += 1; y += e.nbytes;
-          } else {
-            const free = buf_shapes.get(e.key);
-            free.users = e.users;
-            x += 1; y -= free.nbytes;
-            free.x.push(x); free.y.push(free.y.at(-1));
-            temp.delete(e.key);
-            for (const [k, v] of temp) {
-              if (k <= e.key) continue;
-              v.x.push(x, x);
-              v.y.push(v.y.at(-1), v.y.at(-1)-free.nbytes);
-            }
+      const allX = Array.from(new Set(shapes.flatMap(s => s.x))).sort((a,b)=>a-b);
+      const idxs = new Map(allX.map((x,i) => [x, i]));
+      const maxY = new Map(allX.map(x => [x, base0]));
+      // for every [a,b) update the max y at x
+      for (const sh of shapes) {
+        for (let i=0; i<sh.x.length-1; i++) {
+          const startIdx = idxs.get(sh.x[i]), endIdx = idxs.get(sh.x[i+1]);
+          const shapeY = sh.y1[i];
+          for (let k=startIdx; k<endIdx; k++) {
+            const x = allX[k]; maxY.set(x, Math.min(maxY.get(x), shapeY));
           }
         }
-        for (const [num, {dtype, sz, nbytes, y, x:steps, users}] of buf_shapes) {
-          const x = steps.map(s => timestamps[s]);
-          const dur = x.at(-1)-x[0];
-          const arg = { tooltipText:`${dtype}\n${formatUnit(sz)}\n${formatUnit(nbytes, 'B')}\n${formatTime(dur)}`, users, key:`${k}-${bufShapes.length}` };
-          bufShapes.push({ x, y0:y.map(yscale), y1:y.map(y0 => yscale(y0+nbytes)), arg, fillColor:cycleColors(colorScheme.BUFFER, bufShapes.length) });
-        }
-        return bufShapes;
-      };
+      }
+      const sum = {x:[], y0:[], y1:[], fillColor:"#2B1B72"};
+      for (let i=0; i<allX.length-1; i++) {
+        sum.x.push(allX[i], allX[i+1]);
+        const y = maxY.get(allX[i]); sum.y1.push(y, y); sum.y0.push(base0, base0);
+      }
       if (timestamps.length > 0) data.first = data.first == null ? timestamps[0] : Math.min(data.first, timestamps[0]);
       data.tracks.set(k, { shapes:[sum], eventType, visible, offsetY, pcolor:"#c9a8ff", height, peak, scaleFactor:maxheight*4/height,
-                           get views() { return [[sum], buildBufShapes()]; }, valueMap, rowBorderColor });
+                           views:[[sum], shapes], valueMap, rowBorderColor });
       div.style("height", height+padding+"px").style("cursor", "pointer").on("click", (e) => {
         const newFocus = e.currentTarget.id === focusedDevice ? null : e.currentTarget.id;
         let offset = 0;
@@ -512,8 +473,6 @@ async function renderProfiler(path, unit, opts) {
     }
   }
   for (const m of markers) m.label = m.name.split(/(\s+)/).map(st => ({ st, color:m.color, width:ctx.measureText(st).width }));
-  data.pcToShape = new Map([...data.pcToShape].sort((a, b) => a[1].st - b[1].st));
-  if (extData.pcMap != null) data.pcMap = extData.pcMap; setFocus(focusedShape);
   updateProgress(Status.COMPLETE);
   // draw events on a timeline
   const dpr = window.devicePixelRatio || 1;
@@ -539,14 +498,10 @@ async function renderProfiler(path, unit, opts) {
     const visibleX = xscale.range().map(zoomLevel.invertX, zoomLevel).map(xscale.invert, xscale);
     const st = visibleX[0], et = visibleX[1];
     xscale.domain([st, et]);
-    const profilerEl = profiler.node();
-    const visibleYStart = profilerEl.scrollTop-canvasTop + rect(profilerEl).top, visibleYEnd = visibleYStart+profilerEl.clientHeight;
     ctx.textBaseline = "middle";
     // draw shapes
     for (const [k, { shapes, eventType, visible, offsetY, valueMap, pcolor, scolor, rowBorderColor }] of data.tracks) {
       visible.length = 0;
-      const trackHeight = rect(document.getElementById(k)).height;
-      if (offsetY+trackHeight < visibleYStart || offsetY > visibleYEnd) continue;
       const addBorder = scolor != null ? (w) => { if (w > 10) { ctx.strokeStyle = scolor; ctx.stroke(); } } : null;
       for (const e of shapes) {
         if (eventType === EventTypes.BUF) { // generic polygon
@@ -580,7 +535,7 @@ async function renderProfiler(path, unit, opts) {
       }
       // draw row line
       if (rowBorderColor != null) {
-        const y = offsetY+trackHeight-padding/2 - 0.5;
+        const y = offsetY+rect(document.getElementById(k)).height-padding/2 - 0.5;
         drawLine(ctx, [0, canvasWidth], [y, y], { color:rowBorderColor });
       }
     }
@@ -644,7 +599,6 @@ async function renderProfiler(path, unit, opts) {
   document.addEventListener("contextmenu", e => e.ctrlKey && e.preventDefault());
 
   new ResizeObserver(([e]) => e.contentRect.width > 0 && resize()).observe(profiler.node());
-  profiler.on("scroll", () => render(zoomLevel));
 
   function findRectAtPosition(x, y) {
     let track = null;
@@ -859,8 +813,6 @@ async function main() {
   // ** center graph
   const { currentCtx, currentStep, currentRewrite, expandSteps } = state;
   if (currentCtx == -1) return;
-  // always have a new sidebar when view changes
-  metadata.innerHTML = "";
   const ctx = ctxs[currentCtx];
   const step = ctx.steps[currentStep];
   const ckey = step?.query;
@@ -892,6 +844,7 @@ async function main() {
       opts = {heightScale:0.5, hideLabels:true, levelKey:step.name.includes("PKTS") ? (e) => parseInt(e.name.split(" ")[1].split(":")[1]) : null, colorByName:ckey.includes("pkts")};
       return renderProfiler(ckey, "clk", opts);
     }
+    metadata.innerHTML = "";
     ret.metadata?.forEach(m => {
       if (Array.isArray(m)) return metadata.appendChild(tabulate(m.map(({ label, value }) => {
         return [label.trim(), typeof value === "string" ? value : formatUnit(value)];
@@ -1081,15 +1034,13 @@ document.addEventListener("keydown", (event) => {
     if (expandSteps && getSubrewrites(step).length) return step.children[0].click();
     return setState({ expandSteps:!expandSteps });
   }
-  // left and right go through rewrites in a single UOp, in profiler go forward/backward in time
-  if (event.key == "ArrowLeft" || event.key == "ArrowRight") {
+  // left and right go through rewrites in a single UOp
+  if (event.key == "ArrowLeft") {
+    event.preventDefault()
+    return setState({ currentRewrite:Math.max(0, currentRewrite-1) });
+  }
+  if (event.key == "ArrowRight") {
     event.preventDefault()
-    if (profiler.style.display !== "none" && focusedShape != null) {
-      const [t, idx] = focusedShape.split("-");
-      const i = parseInt(idx), last = data.tracks.get(t).shapes.length-1;
-      return setFocus(`${t}-${event.key == "ArrowLeft" ? Math.max(0, i-1) : Math.min(last, i+1)}`);
-    }
-    if (event.key == "ArrowLeft") return setState({ currentRewrite:Math.max(0, currentRewrite-1) });
     const totalRewrites = ret.length-1;
     return setState({ currentRewrite:Math.min(totalRewrites, currentRewrite+1) });
   }

tinygrad/viz/js/worker.js

@@ -54,14 +54,13 @@ const layoutUOp = (g, { graph, change }, opts) => {
       width = Math.max(width, ctx.measureText(line).width);
       height += lineHeight;
     }
-    const callNode = label.startsWith("CALL\n");
-    if (callNode) callCount++;
-    g.setNode(k, {...rectDims(width, height), label, ref, id:k, color, tag, callNode});
+    if (label.startsWith("CALL\n")) callCount++;
+    g.setNode(k, {...rectDims(width, height), label, ref, id:k, color, tag});
     // add edges
     const edgeCounts = {};
     for (const [_, s] of src) edgeCounts[s] = (edgeCounts[s] || 0)+1;
     for (const [port, s] of src) g.setEdge(s, k, { label: edgeCounts[s] > 1 ? {type:"tag", text:edgeCounts[s]} : {type:"port", text:port},
-      ...(callNode && port === 0 && {color:"#a0a1b8"})});
+      ...(label.startsWith("CALL\n") && port === 0 && {color:"#a0a1b8"})});
     if (change?.includes(parseInt(k))) g.setParent(k, "overlay");
   }
   // optionally hide nodes from the layout
@@ -82,7 +81,7 @@ const layoutUOp = (g, { graph, change }, opts) => {
     const disconnected = new Set();
     for (const n of g.nodes()) {
       const node = g.node(n);
-      if (node.callNode && (opts.showCallSrc ? opts.callSrcMask.has(n) : !opts.callSrcMask.has(n))) {
+      if (node.label.startsWith("CALL\n") && (opts.showCallSrc ? opts.callSrcMask.has(n) : !opts.callSrcMask.has(n))) {
         node.collapsed = true;
         for (const pred of (g.predecessors(n) || [])) {
           const edge = g.edge(pred, n);

tinygrad/viz/serve.py

@@ -49,7 +49,7 @@ uops_colors = {Ops.LOAD: "#ffc0c0", Ops.STORE: "#87CEEB", Ops.CONST: "#e0e0e0",
                Ops.INDEX: "#cef263", Ops.WMMA: "#efefc0", Ops.MULTI: "#f6ccff", Ops.INS: "#eec4ff",
                **{x:"#D8F9E4" for x in GroupOp.Movement}, **{x:"#ffffc0" for x in GroupOp.ALU}, Ops.THREEFRY:"#ffff80",
                Ops.BUFFER_VIEW: "#E5EAFF", Ops.BUFFER: "#B0BDFF", Ops.COPY: "#a040a0", Ops.ENCDEC: "#bf71b6",
-               Ops.CALL: "#00B7C8", Ops.PARAM: "#14686F", Ops.SOURCE: "#c0c0c0", Ops.LINEAR: "#7DF4FF", Ops.BINARY: "#404040",
+               Ops.CALL: "#00B7C8", Ops.PARAM: "#14686F", Ops.SOURCE: "#c0c0c0", Ops.LINEAR: "#808080", Ops.BINARY: "#404040",
                Ops.ALLREDUCE: "#ff40a0", Ops.MSELECT: "#d040a0", Ops.MSTACK: "#d040a0", Ops.CONTIGUOUS: "#FFC14D",
                Ops.BUFFERIZE: "#FF991C", Ops.REWRITE_ERROR: "#ff2e2e", Ops.AFTER: "#8A7866", Ops.END: "#524C46"}
 
@@ -128,8 +128,6 @@ def uop_to_json(x:UOp) -> dict[int, dict]:
         label += f"\n({multirange_str(rngs, color=True)})"
       if u._shape is not None:
         label += f"\n{shape_to_str(u.shape)}"
-      if u.op is Ops.CALL:
-        label += f"\n{u.src[0].key.hex()[:8]} {u.src[0].op}"
       if u.op in {Ops.INDEX, Ops.BUFFERIZE}:
         if len(u.toposort()) < 30: label += f"\n{u.render()}"
         ranges: list[UOp] = []
@@ -342,7 +340,7 @@ def sqtt_timeline(data:bytes, lib:bytes, target:str) -> list[ProfileEvent]:
   def add(name:str, p:PacketType, idx=0, width=1, op_name=None, wave=None, info:InstructionInfo|None=None) -> None:
     if hasattr(p, "wave"): wave = p.wave
     rows.setdefault(r:=(f"WAVE:{wave}" if wave is not None else f"{p.__class__.__name__}:0 {name}"))
-    key = TracingKey(f"{op_name if op_name is not None else name} OP:{idx}", ret=f"PC:{info.pc}" if info is not None else None)
+    key = TracingKey(f"{op_name if op_name is not None else name} OP:{idx}", ret=str(info.inst) if info is not None else None)
     ret.append(ProfileRangeEvent(r, key, Decimal(p._time), Decimal(p._time+width)))
   for p, info in map_insts(data, lib, target):
     if len(ret) > getenv("MAX_SQTT_PKTS", 50_000): break
@@ -361,8 +359,7 @@ def sqtt_timeline(data:bytes, lib:bytes, target:str) -> list[ProfileEvent]:
         add(name.replace("_ALT", ""), p, op_name=name)
       if p._time in trace.setdefault(name, set()): raise AssertionError(f"packets overlap in shared resource! {name}")
       trace[name].add(p._time)
-  pc_map = {addr:str(inst) for addr,inst in amd_decode(lib, target).items()}
-  return [ProfilePointEvent(r, "JSON", "pcMap", pc_map, ts=Decimal(0)) for r in rows]+ret
+  return [ProfilePointEvent(r, "start", r, ts=Decimal(0)) for r in rows]+ret
 
 # ** SQTT OCC only unpacks wave start, end time and SIMD location
 
@@ -416,7 +413,6 @@ def get_profile(profile:list[ProfileEvent], sort_fn:Callable[[str], Any]=device_
   # map events per device
   dev_events:dict[str, list[tuple[int, int, float, DevEvent]]] = {}
   markers:list[ProfilePointEvent] = []
-  ext_data:dict[str, Any] = {}
   start_ts:int|None = None
   end_ts:int|None = None
   for ts,en,e in flatten_events(profile):
@@ -424,7 +420,6 @@ def get_profile(profile:list[ProfileEvent], sort_fn:Callable[[str], Any]=device_
     if start_ts is None or st < start_ts: start_ts = st
     if end_ts is None or et > end_ts: end_ts = et
     if isinstance(e, ProfilePointEvent) and e.name == "marker": markers.append(e)
-    if isinstance(e, ProfilePointEvent) and e.name == "JSON": ext_data[e.key] = e.arg
   if start_ts is None: return None
   # return layout of per device events
   layout:dict[str, bytes|None] = {}
@@ -437,8 +432,7 @@ def get_profile(profile:list[ProfileEvent], sort_fn:Callable[[str], Any]=device_
     layout[f"{k} Memory"] = mem_layout(v, start_ts, unwrap(end_ts), peaks, dtype_size, scache)
   sorted_layout = sorted([k for k,v in layout.items() if v is not None], key=sort_fn)
   ret = [b"".join([struct.pack("<B", len(k)), k.encode(), unwrap(layout[k])]) for k in sorted_layout]
-  index = json.dumps({"strings":list(scache), "dtypeSize":dtype_size, "markers":[{"ts":rel_ts(e.ts, start_ts), **e.arg} for e in markers],
-                      **ext_data}).encode()
+  index = json.dumps({"strings":list(scache), "dtypeSize":dtype_size, "markers":[{"ts":rel_ts(e.ts, start_ts), **e.arg} for e in markers]}).encode()
   return struct.pack("<IQII", rel_ts(unwrap(end_ts), start_ts), max(peaks,default=0), len(index), len(ret))+index+b"".join(ret)
 
 # ** PMA counters