print

* allreduce as a function

* allreduce function

* support allreduce function

* LATE_ALLREDUCE

.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/nimlgen/amdcomgr_dylib/releases/latest | \
+        curl -s -H "Authorization: token $GH_TOKEN" curl -s https://api.github.com/repos/tinygrad/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,8 +520,9 @@ 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
-    - 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
+    # 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: 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
 
@@ -616,6 +617,27 @@ 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,6 +244,37 @@ 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
@@ -268,20 +299,6 @@ 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
@@ -644,7 +661,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_sqtt_decoder.py
+        run: sudo PYTHONPATH="." ./extra/sqtt/install_rocprof_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,6 +396,7 @@ 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
+from tinygrad.helpers import getenv, BEAM, WINO, round_up, diskcache_clear, Profiling, profile_marker, DEBUG
 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,11 +1336,13 @@ 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 (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}")
 
+  # pad vocab
+  if (MP := getenv("MP", 1)) > 1: model_params['vocab_size'] = round_up(model_params['vocab_size'], 256 * MP)
+  vocab_mask:Tensor = Tensor.arange(model_params['vocab_size']).reshape(1, 1, -1) >= real_vocab_size
+
   model = Transformer(**model_params, max_context=SEQLEN, jit=False, disable_kv_cache=True)
   params = get_parameters(model)
   # weights are all bfloat16 for now
@@ -1385,8 +1387,10 @@ def train_llama3():
 
   # init grads
   for p in optim.params:
-    p.grad = p.zeros_like().contiguous().realize()
+    p.grad = p.empty_like().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)
 
@@ -1401,15 +1405,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.shard(device, 0)
+      tokens = tokens.to(None).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(-float("inf"), logits).sparse_categorical_crossentropy(tokens[:, 1:])
+    loss = vocab_mask.where(-1e9, 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)
@@ -1417,35 +1421,37 @@ def train_llama3():
 
   @TinyJit
   def optim_step():
-    optim.step()
+    grad_norm = optim.fstep(grads)
     scheduler.step()
 
     for g in grads:
-      g.assign(g.zeros_like())
+      g.assign(g.zeros_like()).realize()
 
     lr = optim.lr
     Tensor.realize(lr, *grads)
 
-    return lr.float().to("CPU")
+    return lr.float().to("CPU"), grad_norm.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.shard(device, 0)
+      tokens = tokens.to(None).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(-float("inf"), logits).sparse_categorical_crossentropy(tokens[:, 1:])
+    loss = vocab_mask.where(-1e9, 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):
-      yield Tensor.randint(bs, SEQLEN + 1, low=0, high=model_params["vocab_size"], dtype=dtypes.int32, device=Device.DEFAULT)
+      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")
 
   def get_train_iter():
     if getenv("FAKEDATA", 0):
@@ -1478,7 +1484,7 @@ def train_llama3():
 
       stopped = False
       losses, data_time, dev_time = [], 0, 0
-      for _ in range(grad_acc):
+      for _ in range(grad_acc if i >= 3 else 1):
         ist = time.perf_counter()
         try: tokens = next(train_iter)
         except StopIteration:
@@ -1491,7 +1497,8 @@ def train_llama3():
       if stopped: break
 
       gt = time.perf_counter()
-      lr = optim_step().item()
+      ret = optim_step()
+      lr, grad_norm = ret[0].item(), ret[1].item()
       et = time.perf_counter()
 
       loss = sum(losses) / len(losses)
@@ -1509,11 +1516,14 @@ 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, {mem_gb:.2f} GB used, {gflops:9.2f} GFLOPS, {mfu:5.2f}% MFU")
+          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())))
 
       if WANDB:
         wandb.log({
-          "lr": lr, "train/loss": loss,
+          "train/loss": loss,
+          "train/lr": lr,
+          "train/grad_norm": grad_norm,
           "train/step_time": step_time,
           "train/gbs_time": gbs_time,
           "train/optim_time": optim_time,
@@ -1550,7 +1560,7 @@ def train_llama3():
       # run eval
       eval_losses = []
       eval_iter = get_eval_iter()
-      tqdm.write(f"evaluating {5760//EVAL_BS} batches of {EVAL_BS} sequences")
+      tqdm.write(f"evaluating {EVAL_SAMPLES//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,40 +7,49 @@ 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).contiguous() for _ in [b1, b2])
+    self.b1_t, self.b2_t = (Tensor.ones((1,), dtype=dtypes.float32, device=self.device, requires_grad=False) 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] = grads[i].to(self.m[i].device)
+      if grads[i].device != self.m[i].device: grads[i].assign(grads[i].to(self.m[i].device))
 
     if self.fused:
-      grads[0] = grads[0] / self.grad_acc
+      grads[0].assign(grads[0] / self.grad_acc)
       total_norm = grads[0].float().square().sum().sqrt()
-      grads[0] = (grads[0] * (self.clip_norm / (total_norm + 1e-6)).clamp(max_=1.0)).cast(grads[0].dtype)
+      grads[0].assign((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] = 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()
+        grads[i].assign(grads[i] / self.grad_acc).realize()
+      total_norm = Tensor.stack(*[g.float().square().sum() for g in grads]).sum().sqrt().contiguous().realize()
       for i in range(len(grads)):
-        grads[i] = (grads[i] * (self.clip_norm / (total_norm + 1e-6)).clamp(max_=1.0)).cast(grads[i].dtype)
+        grads[i].assign((grads[i] * (self.clip_norm / (total_norm + 1e-6)).clamp(max_=1.0)).cast(grads[i].dtype)).realize()
 
     ret = []
     self.b1_t *= self.b1
     self.b2_t *= self.b2
-    for i, (t, g) in enumerate(zip(params, grads)):
+    for i, g in enumerate(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(t.dtype))
-    return ret, [self.b1_t, self.b2_t] + self.m + self.v
+      ret.append((self.lr * up).cast(g.dtype))
+    return ret, [self.b1_t, self.b2_t] + self.m + self.v + [total_norm]
 
   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,6 +5,7 @@ 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
-PYTHONPATH="." extra/viz/cli.py --profile --device "AMD" --top 20
+extra/viz/cli.py --profile --device "AMD" --top 20

examples/openpilot/load_pickle.py

@@ -0,0 +1,16 @@
+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,7 +34,8 @@ class WallTimeEvent:
     self.start = time.monotonic()
     return self
   def __exit__(self, *_):
-    _events[self.event]["wall"].append(time.monotonic() - self.start)
+    self.time = time.monotonic() - self.start
+    _events[self.event]["wall"].append(self.time)
     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, dedup
+from tinygrad.helpers import getenv, all_same, DEBUG
 from extra.gemm.asm.cdna.asm import build_kernel, TILE_M, TILE_N, TILE_K, NUM_WG
 
 # ** CDNA4 assembly gemm
@@ -26,17 +26,25 @@ 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
-atexit.register(lambda: print(f'asm_gemm: {counters["used"]} used, {len(counters["todos"])} not used'))
+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)
 
 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 == 1 and b.uop.axis == 0: K //= len(a.device)
+    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)
     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
@@ -78,6 +86,10 @@ 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)
 
@@ -85,9 +97,16 @@ 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:
-    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)
+    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)
   else:
     out = Tensor.empty(batch, M, N, dtype=a.dtype, device=a.device)
 
@@ -98,4 +117,6 @@ 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)
-  return out.squeeze(0) if squeeze else out
+  out = out.squeeze(0) if squeeze else out
+  if unfold_batch: out = out.reshape(orig_batch, -1, out.shape[-1])
+  return out

extra/sqtt/README.md

@@ -2,30 +2,13 @@
 
 ## Getting SQ Thread Trace
 
-SQTT is implemented on top of normal tinygrad profiling, `VIZ=1 SQTT=1` to get profile pickle with sqtt data embedded in it.
+`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_BUFFER_SIZE=X` to change size of SQTT buffer (per shader engine, 6 SEs on 7900xtx) in megabytes, default 256.
 
-`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
+## Viewing the traces
 
-## 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
-```
+- Web UI: `tinygrad/viz/serve.py`
+- Command line: `python -m tinygrad.renderer.amd.sqtt`

extra/sqtt/active_sqtt_parse.py

@@ -1,152 +0,0 @@
-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

@@ -1,548 +0,0 @@
-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

@@ -0,0 +1,148 @@
+#!/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,6 +9,7 @@ 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_sqtt_decoder.py to install")
+      exc = RuntimeError("Failed to find rocprof-trace-decoder. Run sudo ./extra/sqtt/install_rocprof_decoder.py to install")
       exc.__cause__ = e
   (t:=threading.Thread(target=worker, daemon=True)).start()
   t.join()

extra/thunder/amd/fa.py

@@ -47,19 +47,24 @@ def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False
 
   def grad(dou:UOp, _) -> tuple[None, None, UOp, UOp, UOp]:
     do = Tensor(dou, device=dou.device)
-    dq_in = _sharded_empty((B, H, N, D), xq, axis=shard_axis_t)
-    dq = _sharded_empty_like(xq, axis=shard_axis)
-    dk = _sharded_empty_like(xk, axis=shard_axis)
-    dv = _sharded_empty_like(xv, axis=shard_axis)
+    dq = _sharded_empty((B, H, N, D), xq, axis=shard_axis_t)
+    GROUP_SIZE = H_local // H_KV_local
+    dk_partial = _sharded_empty((B * GROUP_SIZE, N, H_KV, D), xk, axis=shard_axis)
+    dv_partial = _sharded_empty((B * GROUP_SIZE, N, H_KV, D), xv, axis=shard_axis)
 
     # delta_vec = (do * attn).sum(-1, dtype=dtypes.float32).transpose(1, 2).unsqueeze(-2).detach()
     delta_vec = _sharded_empty((B, H, 1, N), xq, dtype=dtypes.float32, axis=shard_axis_t)
-    delta_vec, dq_in = Tensor.custom_kernel(delta_vec, dq_in, attn, do, fxn=functools.partial(custom_fa_backward_pre, device=single_device, arch=arch, B=B_local, N=N, H=H_local, H_KV=H_KV_local, D=D))[:2]
+    delta_vec, dq = Tensor.custom_kernel(delta_vec, dq, attn, do, fxn=functools.partial(custom_fa_backward_pre, device=single_device, arch=arch, B=B_local, N=N, H=H_local, H_KV=H_KV_local, D=D))[:2]
 
-    dq_in, dk, dv = Tensor.custom_kernel(dq_in, dk, dv, do, xq, xk, xv, l_vec, delta_vec, fxn=functools.partial(custom_fa_backward, device=single_device, arch=arch, B=B_local, N=N, H=H_local, H_KV=H_KV_local, D=D))[:3]
+    dq, dk_partial, dv_partial = Tensor.custom_kernel(dq, dk_partial, dv_partial, do, xq, xk, xv, l_vec, delta_vec, fxn=functools.partial(custom_fa_backward, device=single_device, arch=arch, B=B_local, N=N, H=H_local, H_KV=H_KV_local, D=D))[:3]
 
-    # unshuffle dq
-    dq = Tensor.custom_kernel(dq, dq_in, fxn=functools.partial(custom_fa_backward_post, device=single_device, arch=arch, B=B_local, N=N, H=H_local, H_KV=H_KV_local, D=D))[0]
+    # unshuffle dq: atomic_pk_add_bf16_with_warpid creates a shuffled layout within each 16x128 tile
+    # decompose each tile into (j=4, a=2, b=2, d=4, e=4, k=4, c=2) and permute to (e, k, j, a, d, b, c) = standard row-major
+    dq = dq.reshape(B, H, N//16, 4, 2, 2, 4, 4, 4, 2).permute(0, 1, 2, 7, 8, 3, 4, 6, 5, 9).reshape(B, H, N, D).transpose(1, 2)
+
+    # reduce partial dK/dV across GROUP_SIZE query heads
+    dk = dk_partial.reshape(B, GROUP_SIZE, N, H_KV, D).sum(1)
+    dv = dv_partial.reshape(B, GROUP_SIZE, N, H_KV, D).sum(1)
 
     return None, None, dq.uop, dk.uop, dv.uop
 
@@ -89,7 +94,6 @@ def custom_fa_forward(o:UOp, l_vec:UOp, q:UOp, k:UOp, v:UOp, device:str, arch:st
                   arg=KernelInfo(name="custom_fa_forward", estimates=estimates))
 
   lib = HIPCCCompiler(arch, compile_args).compile_cached(code)
-
   lib = bytearray(lib)
   rodata_off = next(sh.header.sh_offset for sh in elf_loader(bytes(lib))[1] if sh.name == ".rodata")
   struct.pack_into('<I', lib, rodata_off, 160000)
@@ -120,7 +124,6 @@ def custom_fa_backward_pre(delta_vec:UOp, dq:UOp, o:UOp, do:UOp, device:str, arc
                   arg=KernelInfo(name="custom_fa_backward_pre", estimates=estimates))
 
   lib = HIPCCCompiler(arch, compile_args).compile_cached(code)
-
   lib = bytearray(lib)
   rodata_off = next(sh.header.sh_offset for sh in elf_loader(bytes(lib))[1] if sh.name == ".rodata")
   struct.pack_into('<I', lib, rodata_off, 160000)
@@ -138,7 +141,7 @@ def custom_fa_backward(dq:UOp, dk:UOp, dv:UOp, do:UOp, q:UOp, k:UOp, v:UOp, l_ve
   BLOCK_SIZE_KV = 256
   NUM_WARPS = 4
   NUM_THREADS = 64 * NUM_WARPS
-  gsz = (H_KV, N // BLOCK_SIZE_KV, B)
+  gsz = (H, N // BLOCK_SIZE_KV, B)
   lsz = (NUM_THREADS, 1, 1)
   threadIdx_x = UOp.special(lsz[0], "lidx0")
   blockIdx_x, blockIdx_y, blockIdx_z = UOp.special(gsz[0], "gidx0"), UOp.special(gsz[1], "gidx1"), UOp.special(gsz[2], "gidx2")
@@ -151,7 +154,6 @@ def custom_fa_backward(dq:UOp, dk:UOp, dv:UOp, do:UOp, q:UOp, k:UOp, v:UOp, l_ve
                   arg=KernelInfo(name="custom_fa_backward", estimates=estimates))
 
   lib = HIPCCCompiler(arch, compile_args).compile_cached(code)
-
   lib = bytearray(lib)
   rodata_off = next(sh.header.sh_offset for sh in elf_loader(bytes(lib))[1] if sh.name == ".rodata")
   struct.pack_into('<I', lib, rodata_off, 160000)
@@ -182,7 +184,6 @@ def custom_fa_backward_post(dq_out:UOp, dq_in:UOp, device:str, arch:str, B:int,
                   arg=KernelInfo(name="custom_fa_backward_post", estimates=estimates))
 
   lib = HIPCCCompiler(arch, compile_args).compile_cached(code)
-
   lib = bytearray(lib)
   rodata_off = next(sh.header.sh_offset for sh in elf_loader(bytes(lib))[1] if sh.name == ".rodata")
   struct.pack_into('<I', lib, rodata_off, 160000)

extra/thunder/amd/fa_bwd_causal.cpp

@@ -37,7 +37,7 @@ using namespace kittens;
 using _gl_QdO  = gl<bf16, ATTN_B, ATTN_N, ATTN_H, ATTN_D>;
 using _gl_KV   = gl<bf16, ATTN_B, ATTN_N, ATTN_H_KV, ATTN_D>;
 using _gl_dQ   = gl<bf16, ATTN_B, ATTN_H, ATTN_N, ATTN_D>;
-using _gl_dKV  = gl<bf16, ATTN_B, ATTN_N, ATTN_H_KV, ATTN_D>;
+using _gl_dKV  = gl<bf16, ATTN_B * GROUP_SIZE, ATTN_N, ATTN_H_KV, ATTN_D>;
 using _gl_Lvec = gl<float, ATTN_B, ATTN_H, 1, ATTN_N>;
 
 template<int D> struct attn_bwd_combined_globals {
@@ -47,7 +47,7 @@ template<int D> struct attn_bwd_combined_globals {
   _gl_dQ dQg;
   _gl_dKV dKg, dVg;
   _gl_Lvec L_vec, delta_vec;
-  dim3 grid() { return dim3(ATTN_H_KV, (ATTN_N / BLOCK_SIZE_KV), ATTN_B); }
+  dim3 grid() { return dim3(ATTN_H, (ATTN_N / BLOCK_SIZE_KV), ATTN_B); }
   dim3 block() { return dim3(NUM_THREADS); }
   size_t dynamic_shared_memory() { return MAX_SHARED_MEMORY; }
 };
@@ -55,10 +55,12 @@ template<int D> struct attn_bwd_combined_globals {
 template<int D> __launch_bounds__(NUM_THREADS, 1)
 __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr, bf16 *dO_ptr, bf16 *Q_ptr, bf16 *K_ptr, bf16 *V_ptr, float *L_vec_ptr, float *delta_vec_ptr) {
 
-  const int kv_head_idx = blockIdx.x;  // This is the KV head index
+  const int q_head_idx_fixed = blockIdx.x;  // This is the query head index [0, ATTN_H)
+  const int kv_head_idx = q_head_idx_fixed / GROUP_SIZE;
+  const int q_head_in_group = q_head_idx_fixed % GROUP_SIZE;
   const int seq_idx = blockIdx.y;
   const int batch_idx = blockIdx.z;
-  const int first_q_head = kv_head_idx * GROUP_SIZE;
+  const int first_q_head = q_head_idx_fixed;
 
   const int warpid = kittens::warpid();
   const int j = seq_idx * NUM_WARPS + warpid;
@@ -70,7 +72,7 @@ __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr
   // first Q step that can overlap this K_span:
   const int first_step = max(0, k_start_min / STEP_QO);
   const int num_steps_per_head = total_steps_per_head - first_step;
-  const int num_steps = num_steps_per_head * GROUP_SIZE;
+  const int num_steps = num_steps_per_head;
   const int k_pos = j * WARP_SIZE_KV;
 
   constexpr float L_SCALE_FACTOR = 1.44269504089f;
@@ -3355,14 +3357,14 @@ __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr
     }
   }
 
-  store<1>(g.dVg, dV_j, {batch_idx, 0, kv_head_idx, 0}, {0, j, 0, 0});
+  store<1>(g.dVg, dV_j, {batch_idx * GROUP_SIZE + q_head_in_group, 0, kv_head_idx, 0}, {0, j, 0, 0});
   __builtin_amdgcn_s_waitcnt(0);
   __builtin_amdgcn_s_barrier();
 
   // We first copy dV_j_T from accumulator GPRs to vector GPRs and then perform the store
   accvgpr_read(dV_j_T, dK_j_T);
   mul(dV_j_T, dV_j_T, dP_SCALE_FACTOR);
-  store<1>(g.dKg, dV_j, {batch_idx, 0, kv_head_idx, 0}, {0, j, 0, 0});
+  store<1>(g.dKg, dV_j, {batch_idx * GROUP_SIZE + q_head_in_group, 0, kv_head_idx, 0}, {0, j, 0, 0});
 
   // Write out final dQ_i slice
   mul(dQ_i_T, dQ_i_T, dP_SCALE_FACTOR);

extra/torch_backend/test_kernel_fusion.py

@@ -1,7 +1,6 @@
 # simple tests
 import unittest
 import torch
-import warnings
 from tinygrad.helpers import getenv, GlobalCounters
 if getenv("TINY_BACKEND2"):
   import extra.torch_backend.backend2
@@ -18,9 +17,7 @@ class TestKernelFusionRegression(unittest.TestCase):
     torch.manual_seed(42)
     GlobalCounters.reset()
     fn().detach().cpu().numpy()
-    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}")
+    self.assertEqual(GlobalCounters.kernel_count, expected_kernels)
 
   def test_elementwise_fusion(self):
     def fn():
@@ -34,7 +31,7 @@ class TestKernelFusionRegression(unittest.TestCase):
       conv = torch.nn.Conv2d(3, 16, 3, padding=1).to(device)
       with torch.no_grad():
         return torch.nn.functional.relu(conv(x))
-    self._check_kernel_count(fn, 8)
+    self._check_kernel_count(fn, 6)
 
   def test_batchnorm_fusion(self):
     def fn():
@@ -44,7 +41,7 @@ class TestKernelFusionRegression(unittest.TestCase):
       bn.eval()
       with torch.no_grad():
         return torch.nn.functional.relu(bn(conv(x)))
-    self._check_kernel_count(fn, 16)
+    self._check_kernel_count(fn, 10)
 
   def test_reduce_fusion(self):
     def fn():
@@ -92,7 +89,7 @@ class TestKernelFusionRegression(unittest.TestCase):
         out = bn(conv(x))
         out += identity
         return torch.nn.functional.relu(out)
-    self._check_kernel_count(fn, 17)
+    self._check_kernel_count(fn, 12)
 
   def test_multiple_inplace_ops_fusion(self):
     def fn():
@@ -117,7 +114,7 @@ class TestKernelFusionRegression(unittest.TestCase):
       bn.train()
       with torch.no_grad():
         return bn(x)
-    self._check_kernel_count(fn, 10)
+    self._check_kernel_count(fn, 8)
 
   # this is a minimal extra/other_mnist/beautiful_mnist_torch.py to cover fusion for training with optimizer
   def test_mnist_training_fusion(self):
@@ -138,7 +135,7 @@ class TestKernelFusionRegression(unittest.TestCase):
       loss.backward()
       optimizer.step()
       return loss
-    self._check_kernel_count(fn, 28)
+    self._check_kernel_count(fn, 24)
 
 if __name__ == "__main__":
   unittest.main()

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

@@ -0,0 +1,17 @@
+<?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

@@ -0,0 +1,33 @@
+#!/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 `PYTHONPATH=. extra/viz/cli.py` to explore the saved trace files.
+After running with VIZ=-1, use `extra/viz/cli.py` to explore the saved trace files.
 
 ## Inspect runtime profiling
 
-Use `PYTHONPATH=. extra/viz/cli.py --profile` to list all traced devices.
+Use `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 `PYTHONPATH=. extra/viz/cli.py --rewrites` to list all traced kernels.
+Use `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,44 +1,73 @@
 #!/usr/bin/env python3
 import os
 os.environ["VIZ"] = "0"
-import argparse, pathlib
+import argparse, pathlib, sys, struct, json
 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
-from test.null.test_viz import load_profile
+
+# ** generic helpers
 
 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("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 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 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)",
@@ -46,14 +75,14 @@ if __name__ == "__main__":
   args = parser.parse_args()
   if not args.profile and not args.rewrites:
     parser.print_help()
-    exit(0)
+    sys.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 = load_profile(viz.load_pickle(args.profile_path, default=[]))
+    profile = decode_profile(viz.get_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():
@@ -63,7 +92,7 @@ if __name__ == "__main__":
       for e in v.get("events", []):
         et = e["dur"]*1e-6
         if args.kernel is not None:
-          if ansistrip(e["name"]) == args.kernel and n < 10:
+          if optional_eq(e, 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', ' | ')+"  ")
@@ -81,7 +110,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"))
-    exit(0)
+    sys.exit(0)
 
   for k in viz.ctxs:
     if not optional_eq(k, args.kernel): continue

test/amd/disasm.py

@@ -324,6 +324,12 @@ 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)
@@ -331,9 +337,10 @@ 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
+  # Global/scratch uses 13-bit signed offset (RDNA3/CDNA), 24-bit signed offset (RDNA4)
   offset = inst.ioffset if r4 else inst.offset  # type: ignore[attr-defined]
-  if seg != 'flat':
+  if r4: off_val = offset if offset < (1 << 23) else offset - (1 << 24)  # sign extend 24-bit
+  elif seg != 'flat':
     if cdna:
       # CDNA: bit 12 is sign bit but not in offset field
       raw = int.from_bytes(inst.to_bytes(), 'little')
@@ -348,7 +355,9 @@ 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: mods = f"{off_s}{' scope' if inst.scope else ''}{' th' if inst.th 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 "")
   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"
@@ -357,7 +366,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.data if 'store' in name else inst.vdst)}{saddr_s}{mods}"
+  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}"
    # 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]
@@ -372,7 +381,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.glc  # type: ignore[attr-defined]
+    glc_or_sc0 = inst.sc0 if cdna else (inst.th & 1 if r4 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,6 +104,34 @@ 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_ds.s', 'gfx12_asm_smem.s', 'gfx12_asm_vflat.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, InstOpRDNA4.OTHER_VMEM_STORE}
+OTHER_SIMD_OPS_RDNA4 = {InstOpRDNA4.OTHER_VMEM}
 
 # ═══════════════════════════════════════════════════════════════════════════════
 # ROCPROF DECODER

test/amd/test_sqtt_tables.py

@@ -1,183 +0,0 @@
-"""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
+from tinygrad.helpers import DEBUG, OSX
 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):
+def rocprof_inst_traces_match(sqtt, prg, target, pass_rocprof_err=False):
   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):
 
   passed_insts = 0
   for pkt, info in map_insts(sqtt.blob, prg.lib, target):
-    if DEBUG >= 2: print_packets([pkt])
+    if DEBUG >= 2: print_packets([(pkt, info)])
     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, f"pc mismatch {ref_pc}:{disasm_map[rocprof_inst.pc]} != {info.pc}:{disasm(info.inst)}"
+    assert ref_pc == info.pc or pass_rocprof_err, 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,7 +67,9 @@ 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):
-          passed_insts, n_waves, n_units = rocprof_inst_traces_match(event, kern_events[event.kern], target)
+          # 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)
           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,6 +47,18 @@ 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")
@@ -60,6 +72,14 @@ 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):
@@ -101,6 +121,20 @@ 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,6 +27,11 @@ 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_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_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_allreduce_all2all(self):
     with Context(ALL2ALL=2):

test/backend/test_schedule.py

@@ -795,6 +795,38 @@ 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,6 +205,20 @@ 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]])
@@ -279,17 +293,43 @@ 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.rand(8, 8)
-    x = Tensor.rand(8, requires_grad=True)
-    with self.assertRaises(NotImplementedError):
-      z[:3] = x
+    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])
 
 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(AssertionError, "bind mismatch on"):
+    with self.assertRaisesRegex(RuntimeError, "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
+from tinygrad.helpers import Context, getenv
 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"}, "only NV, AMD, CUDA")
+  @unittest.skipIf(Device.DEFAULT not in {"CUDA", "NV", "AMD"} or getenv("MOCKGPU"), "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,6 +69,58 @@ 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,6 +136,30 @@ 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

@@ -0,0 +1,31 @@
+#!/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

@@ -0,0 +1,55 @@
+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,13 +4,19 @@
 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
+import unittest, re, importlib
 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
 
-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}
+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
   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"
 
@@ -21,6 +27,10 @@ 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
@@ -33,7 +43,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)))
+    prg = AMDProgram(self.dev, "test", self.compiler.compile(assemble(code, is_cdna=self.is_cdna)))
     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]):
@@ -57,32 +67,32 @@ class TestOutOfBoundsMemoryAccess(TestGPUCrash):
 
   def test_global_load_null_ptr(self):
     """Global load from NULL pointer."""
-    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()]
+    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()]
     self._assert_gpu_fault(lambda: self._run_insts(insts))
 
   def test_global_store_null_ptr(self):
     """Global store to NULL pointer."""
-    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()]
+    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()]
     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 = [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()]
+    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()]
     self._assert_gpu_fault(lambda: self._run_insts(insts))
 
   def test_global_store_unmapped_high_address(self):
     """Global store to high unmapped address."""
-    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()]
+    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()]
     self._assert_gpu_fault(lambda: self._run_insts(insts))
 
   def test_global_atomic_unmapped(self):
     """Atomic operation on unmapped memory."""
-    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()]
+    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()]
     self._assert_gpu_fault(lambda: self._run_insts(insts))
 
 
@@ -91,14 +101,14 @@ class TestSMEMFaults(TestGPUCrash):
 
   def test_smem_load_null(self):
     """SMEM load from NULL base."""
-    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()]
+    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()]
     self._assert_gpu_fault(lambda: self._run_insts(insts))
 
   def test_smem_load_unmapped(self):
     """SMEM load from unmapped address."""
-    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()]
+    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()]
     self._assert_gpu_fault(lambda: self._run_insts(insts))
 
 
@@ -107,20 +117,20 @@ class TestFlatMemoryFaults(TestGPUCrash):
 
   def test_flat_load_null(self):
     """FLAT load from NULL address."""
-    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()]
+    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()]
     self._assert_gpu_fault(lambda: self._run_insts(insts))
 
   def test_flat_store_null(self):
     """FLAT store to NULL address."""
-    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()]
+    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()]
     self._assert_gpu_fault(lambda: self._run_insts(insts))
 
   def test_flat_atomic_null(self):
     """FLAT atomic on NULL address."""
-    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()]
+    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()]
     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): pass
+  def __call__(self, *bufs, global_size, local_size, vals=(), wait=False, **kw): 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) if op == '==' else left.ne(right)
+      case '==': return left.eq(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)}
+      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))}
         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)
+    self._validate(repo_id, model_file, custom_inputs, atol=1e-3)
 
 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.assertEqual(GlobalCounters.kernel_count, 1)
+    self.assertLessEqual(GlobalCounters.kernel_count, 1)  # 0 if BUFFER_VIEW (zero-copy), 1 otherwise
 
 if __name__ == '__main__':
   unittest.main()

test/null/test_multitensor.py

@@ -1,10 +1,12 @@
 import gc, unittest
 from tinygrad import Tensor, GlobalCounters, dtypes
+from tinygrad.engine.jit import TinyJit
 
 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()
@@ -12,6 +14,11 @@ 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()
@@ -59,6 +66,33 @@ 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 = {}
@@ -74,6 +108,38 @@ class TestMultiRamUsage(unittest.TestCase):
   def test_matmul_half(self): self._test_matmul_half(dev_count=2)
   def test_matmul_half_alt(self): self._test_matmul_half(dev_count=4)
 
+  @unittest.expectedFailure
+  def test_multi_layer_allreduce(self):
+    N = 32
+    devices_2 = ("NULL:1", "NULL:2")
+
+    def make_inp():
+      x = Tensor.zeros(N, N).contiguous().shard(devices_2, axis=None).realize()
+      w1 = Tensor.zeros(N, N).contiguous().shard(devices_2, axis=1).realize()
+      w2 = Tensor.zeros(N, N).contiguous().shard(devices_2, axis=0).realize()
+      return x, w1, w2
+
+    def run_layers(n_layers):
+      GlobalCounters.reset()
+
+      @TinyJit
+      def f(x, w1, w2):
+        for _ in range(n_layers):
+          x = (x @ w1 @ w2)
+        return x.contiguous()
+
+      for _ in range(3):
+        a = make_inp()
+        r = f(*a)
+        del a, r
+
+      gc.collect()
+      return GlobalCounters.mem_used
+
+    mem_2 = run_layers(2)
+    mem_4 = run_layers(4)
+    self.assertEqual(mem_2, mem_4, f"graph memory should not grow with layers: 2 layers={mem_2}, 4 layers={mem_4}")
+
 class TestMultiAxis(unittest.TestCase):
   def test_reshape_shard_invalid(self):
     devices = ("NULL:0", "NULL:1")

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, 1)
+    check_schedule(b, 0)  # contiguous shrink of a realized buffer is a zero-copy BUFFER_VIEW
 
   def test_double_contiguous_realizes_once(self):
     a = Tensor.empty(4, 1)
@@ -234,6 +234,18 @@ 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)
@@ -1104,6 +1116,7 @@ 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))
@@ -1157,5 +1170,23 @@ 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, json, struct, sys
+import unittest, decimal, sys
 from dataclasses import dataclass
 from typing import Generator
 
@@ -357,41 +357,9 @@ class TestVizIntegration(BaseTestViz):
 
 from tinygrad.device import ProfileDeviceEvent, ProfileGraphEvent, ProfileGraphEntry
 from tinygrad.viz.serve import get_profile
+from extra.viz.cli import decode_profile
 
-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}
+def load_profile(lst:list[ProfileEvent]) -> dict: return decode_profile(get_profile(lst))
 
 class TestVizProfiler(BaseTestViz):
   def test_transfer_uses_copy_device(self):

test/unit/test_assign.py

@@ -29,6 +29,7 @@ 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()
@@ -269,6 +270,16 @@ class TestAssign(unittest.TestCase):
     out = attn.cache_k.flatten().numpy()
     np.testing.assert_allclose(out, [1.,1.,1.,1.,1.,1.,0.,0.,1.,1.,1.,1.,1.,1.,0.,0.])
 
+  def test_assign_after(self):
+    t = Tensor.zeros(10).contiguous().realize()
+    t.uop = t.uop.after(t.uop.assign((t+1).uop))
+    np.testing.assert_allclose(t.numpy(), [1.,1.,1.,1.,1.,1.,1.,1.,1.,1.])
+
+  def test_assign_after_partial(self):
+    t = Tensor.zeros(10).contiguous().realize()
+    t.uop = t.uop.after(t[:5].uop.assign(Tensor.ones(5).uop))
+    np.testing.assert_allclose(t.numpy(), [1.,1.,1.,1.,1.,0.,0.,0.,0.,0.])
+
   def test_assign_contiguous(self):
     b = Tensor.arange(16).reshape(4,4).contiguous().realize()
     a = (Tensor.arange(16).reshape(4,4).contiguous().realize() + 1)
@@ -598,8 +609,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()
-    # 2 kernels for first assign + 3 per remaining assign (matmul, contiguous, assign) + 1 final read = 3*N
-    self.assertEqual(GlobalCounters.kernel_count, 3*N)
+    # 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)
 
 
 class TestAssignOrdering(unittest.TestCase):
@@ -676,16 +687,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].contiguous()  # lazy - not captured yet
-    right = buf[4:8].contiguous()  # lazy - not captured yet
+    left = buf[0:4].clone()  # lazy - not captured yet
+    right = buf[4:8].clone()  # 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].contiguous().realize()
-    right = buf[4:8].contiguous().realize()
+    left = buf[0:4].clone().realize()
+    right = buf[4:8].clone().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])
@@ -756,13 +767,12 @@ class TestAssignOrdering(unittest.TestCase):
     np.testing.assert_equal(b.numpy(), [1, 2, 3, 4])
 
   def test_variable_slice_ordering(self):
-    """Variable-indexed slices - tests symbolic dependency tracking."""
+    """Variable-indexed slices - conflicting variable binds in same schedule are rejected."""
     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)
-    self.assertEqual(buf[0:1, :].sum().item(), 4)
-    self.assertEqual(buf[1:2, :].sum().item(), 8)
+    with self.assertRaises(RuntimeError): buf[0:1, :].sum().item()
 
   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
+from tinygrad import Tensor, function
 from tinygrad.dtype import dtypes
 from tinygrad.uop.ops import UOp
 
@@ -100,5 +100,43 @@ 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

@@ -2,6 +2,7 @@ import numpy as np
 import unittest
 from tinygrad.function import function
 from tinygrad import Tensor
+from tinygrad.uop.ops import UOp
 
 class TestFunction(unittest.TestCase):
   def test_simple(self):
@@ -69,6 +70,14 @@ 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])
@@ -102,7 +111,6 @@ class TestFunction(unittest.TestCase):
     np.testing.assert_allclose(w.grad.numpy(), [4., 5., 6.])
 
   def test_symbolic_index(self):
-    from tinygrad.uop.ops import UOp
     table = Tensor([10,20,30,40]).contiguous().realize()
     @function
     def f(x:Tensor, start_pos:int|UOp) -> Tensor:
@@ -111,6 +119,14 @@ class TestFunction(unittest.TestCase):
     v = UOp.variable("start_pos", 0, 3)
     np.testing.assert_equal(f(Tensor([1,2,3]), v.bind(0)).numpy(), [11,12,13])
 
+  def test_symbolic_shape_input(self):
+    table = Tensor([10,20,30,40]).contiguous().realize()
+    @function
+    def f(x:Tensor) -> Tensor: return x * 2
+    sz = UOp.variable("sz", 1, 3)
+    slic = table[:sz.bind(2)]
+    np.testing.assert_equal(f(slic)[:2].numpy(), [20,40])
+
   def test_nested_calls(self):
     w = Tensor([10., 20., 30.])
     @function
@@ -121,6 +137,31 @@ 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
@@ -185,5 +226,114 @@ 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, unittest
+import os, struct, 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,5 +120,41 @@ 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,6 +68,14 @@ 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,8 +179,6 @@ 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)
@@ -190,7 +188,6 @@ 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,6 +83,15 @@ 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

@@ -116,6 +116,7 @@ class TransformerBlock:
       self.ffn_up      = nn.Linear(dim, hidden_dim, bias=False)
       self.ffn_down    = nn.Linear(hidden_dim, dim, bias=False)
 
+  @function
   def _attention(self, x:Tensor, start_pos:int|UOp) -> Tensor:
     x_norm = self.attn_norm(x)                       # (B,T,D)
     q, k, v = self.attn_q(x_norm), self.attn_k(x_norm), self.attn_v(x_norm)
@@ -131,18 +132,25 @@ class TransformerBlock:
     q = apply_rope(q, freqs_cis)
     k = apply_rope(k, freqs_cis)
 
-    self.cache_kv[:, :, :, start_pos:start_pos+T, :].assign(Tensor.stack(k, v))
-    k = self.cache_kv[0, :, :, 0:start_pos+T, :]
-    v = self.cache_kv[1, :, :, 0:start_pos+T, :]
+    # TODO: fix assign to behave like this
+    assigned_kv = self.cache_kv.uop.after(self.cache_kv[:, :, :, start_pos:start_pos+T, :].uop.assign(Tensor.stack(k, v).contiguous().uop))
+    tensor_assigned_kv = Tensor(assigned_kv, device=assigned_kv.device)
+    k = tensor_assigned_kv[0, :, :, 0:start_pos+T, :]
+    v = tensor_assigned_kv[1, :, :, 0:start_pos+T, :]
+
+    #self.cache_kv[:, :, :, start_pos:start_pos+T, :].assign(Tensor.stack(k, v))
+    #k = self.cache_kv[0, :, :, 0:start_pos+T, :]
+    #v = self.cache_kv[1, :, :, 0:start_pos+T, :]
 
     # NOTE: this mask is causal_lower_right, not the causal_upper_left generated by is_casual = True
-    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
+    # 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
     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
+  @function(precompile=bool(getenv("PRECOMPILE", 0)))
   def _feed_forward(self, h: Tensor) -> Tensor:
     h_norm = self.ffn_norm(h)
     if hasattr(self, 'ffn_gate_exps'):
@@ -182,9 +190,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=True) -> tuple[Transformer, dict]:
+  def from_gguf(gguf:Tensor, max_context:int|None=None, realize=bool(getenv("REALIZE", 1))) -> tuple[Transformer, dict]:
     # TODO: remove the need for copy to default device
-    kv, state_dict = nn.state.gguf_load(gguf.to(None))
+    kv, state_dict = nn.state.gguf_load(gguf.to(None).realize())
 
     # 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()}
@@ -212,8 +220,9 @@ 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
-    for s in (params:=nn.state.get_parameters(model)): s.replace(s.contiguous())
-    if realize: Tensor.realize(*params)
+    if realize:
+      for s in (params:=nn.state.get_parameters(model)): s.replace(s.contiguous())
+      Tensor.realize(*params)
     return model, kv
 
   def generate(self, tokens:list[int], start_pos=0):
@@ -255,7 +264,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.preventDefault(); send() } }
+  input.onkeydown = (e) => { if (e.key === 'Enter' && !e.shiftKey && !e.isComposing) { e.preventDefault(); send() } }
   const msgs = [];
   async function send() {
     if (!input.value.trim()) return;
@@ -329,37 +338,42 @@ class Handler(HTTPRequestHandler):
 
 if __name__ == "__main__":
   parser = argparse.ArgumentParser()
-  parser.add_argument("--model", choices=list(models.keys()), default=list(models.keys())[0], help="Model choice")
+  parser.add_argument("--model", "-m", 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
-  model, kv = Transformer.from_gguf(Tensor.from_url(models[args.model]), args.max_context)
-  if DEBUG >= 1: print(f"using model {args.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
 
-  # 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)
+  # TODO: why this is required to free the RAM of the GGUF copy?
+  import gc
+  gc.collect()
 
   # 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, field
+from dataclasses import dataclass
 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
@@ -189,7 +189,10 @@ def _do_image_fixup(dt:ImageDType, idx:UOp) -> tuple[UOp, UOp, int, int]:
   h, w = dt.shape[0], dt.shape[1]
   if IMAGE == 1 and valid is not None:
     h, w = max(ImageDType.valid_dims(dt), key=lambda hw:
-               (len(_drop_valid_stmts(valid, idx:=uop_given_valid(valid, UOp.vectorize((x//4)%hw[1], x//(4*hw[1]))), *hw)), -len(idx.backward_slice)))
+                # maximize number of valids removed
+               (len(_drop_valid_stmts(valid, idx:=uop_given_valid(valid, UOp.vectorize((x//4)%hw[1], x//(4*hw[1]))), *hw)),
+                # and minimize idx complexity (number of nodes)
+                -len(idx.simplify().backward_slice)))
     buf = buf.replace(dtype=(dtypes.imageh if dt.itemsize == 2 else dtypes.imagef)((h, w, 4), w * 4 * dt.itemsize))
   oidx = UOp(Ops.VECTORIZE, dtypes.index.vec(2), ((x // 4) % w, (x // (4*w))))
   return x, idx.replace(src=(buf, oidx.valid(valid))), w, h
@@ -308,8 +311,6 @@ 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
@@ -335,13 +336,15 @@ 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)
-  ctx.range_to_ends.setdefault(reduce_range, []).append(end)
+  end = acc.index(UOp.const(dtypes.int, 0)).store(ret).end(*reduce_range).rtag("mergeable")
   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
-  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}
+  # 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}
   return sink.substitute(subs) if subs else None
 
 pm_reduce = PatternMatcher([

tinygrad/codegen/opt/search.py

@@ -36,7 +36,8 @@ 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") -> list[float]:
+                  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
   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)
@@ -50,7 +51,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))*factor)
+    tms.append(unwrap(car(input_bufs, var_vals, wait=True, timeout=timeout))*factor)
     if early_stop is not None and early_stop < min(tms): break
   return tms
 
@@ -161,7 +162,8 @@ 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'))
+                                 allow_test_size=allow_test_size, clear_l2=hasattr(dev, 'invalidate_caches'),
+                                 dev_timeout=getenv("BEAM_DEV_TIMEOUT", 1))
         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,6 +141,7 @@ 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):
@@ -149,6 +150,7 @@ 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, identity_element, track_rewrites
-from tinygrad.dtype import ImageDType
-from tinygrad.helpers import prod, DEBUG, argsort, VIZ, pluralize
+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, VIZ, pluralize
 
 @dataclass
 class AllocCtx:
@@ -25,7 +25,9 @@ def disk_copy_is_buffer(ctx:AllocCtx, u:UOp):
   if from_creation: return tag_uop(ctx, u)
 
 def apply_after(ctx:AllocCtx, u:UOp):
-  ctx.buffer_map[u] = u.src[0]
+  base = u.src[0]
+  while base.op is Ops.AFTER: base = base.src[0]
+  ctx.buffer_map[u] = base
 
 # CONTIGUOUS and ASSIGN + parents are the only nodes that get updated
 add_tags = PatternMatcher([
@@ -38,11 +40,7 @@ add_tags = PatternMatcher([
   (UPat(GroupOp.All, name="x"), lambda ctx,x: tag_uop(ctx,x) if x in ctx.bases else None),
 ])
 
-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)
+def _buffer_like(u:UOp) -> UOp:
   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:
@@ -50,45 +48,72 @@ def replace_contig_with_assign(u: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.assign(u.src[0]).rtag(u.tag)
+  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)
 
 def replace_assign_with_contig(u:UOp):
   assigned_to = u
-  while assigned_to.op in {Ops.ASSIGN, Ops.BITCAST}: assigned_to = assigned_to.src[0].base
+  while assigned_to.op in {Ops.ASSIGN, Ops.BITCAST, Ops.AFTER}: 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):
-  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[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([
-  # 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),
+  # 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),
+
   # 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
-  (UPat(Ops.CONTIGUOUS, src=(UPat(Ops.ASSIGN, name="a"),), name="c"), lambda a,c: a.replace(tag=a.tag+c.tag)),
+  # 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),
   # 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
+  # remove DETACH/CONTIGUOUS_BACKWARD (allows more contiguous removal)
   (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):
@@ -121,6 +146,8 @@ 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),
 ])
@@ -138,10 +165,10 @@ def transform_to_call(big_sink:UOp) -> tuple[UOp, dict[UOp, UOp]]:
   big_sink = graph_rewrite(big_sink, add_tags, ctx=ctx, bottom_up=True, name="number the uops")
 
   # here we can break the tensor graph. this is the only place you need to maintain numbered tags
-  big_sink = graph_rewrite(big_sink, pm_early_transform_tensor_graph, ctx={}, name="early transform tensor graph")
+  big_sink = graph_rewrite(big_sink, pm_early_transform_tensor_graph, name="early transform tensor graph")
 
   # 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, name="replace bufs").call(*ctx.replacements)
+  ret = graph_rewrite(UOp.sink(*ctx.assigns), pm_replace_buf, ctx=ctx, bottom_up=True, 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) -> float|None:
+  def __call__(self, rawbufs:list[Buffer], var_vals:dict[str, int]|None=None, wait=False, timeout:int|None=None) -> 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)
+                     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)
 
 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,10 +2,9 @@ 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
+from tinygrad.helpers import DEBUG, cpu_profile, TracingKey, SPEC, pluralize, SCACHE, BASEDIR, flatten
 from tinygrad.engine.realize import ExecItem
 
 # **** schedule linearizer
@@ -23,7 +22,7 @@ def create_schedule(sched_sink:UOp) -> UOp:
     for u in sched_sink.toposort(gate_kernel_sink):
       if u.op is not Ops.AFTER: continue
       k = u.src[1]
-      assert k.op in {Ops.CALL, Ops.END, Ops.LINEAR}, f"AFTER src[1] should be CALL or END, not {k.op}"
+      assert k.op in {Ops.CALL, Ops.END}, f"AFTER src[1] should be 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:]
@@ -92,27 +91,31 @@ 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].rtag() if x.tag is None else None),
+  (UPat(Ops.PARAM, name="x"), lambda ctx,x: ctx[1][x.arg]),
   # create new BUFFERs for LUNIQUE BUFFERs from rangeify
   (UPat(Ops.BUFFER, src=(UPat(Ops.LUNIQUE), UPat(Ops.DEVICE)), name="b"), create_new_buffer),
 ])
 
-# the AFTER structure is already in LINEAR
-pm_collapse_after = PatternMatcher([
-  (UPat(Ops.AFTER, name="x"), lambda x: x.src[0])
+pm_resolve_linear_call = PatternMatcher([
+  # call LINEAR is resolved here
+  (UPat(Ops.CALL, src=(UPat(Ops.LINEAR),), name="linear_call", allow_any_len=True), lambda linear_call:
+   graph_rewrite(linear_call.src[0], pm_post_sched_cache, ctx=({}, linear_call.src[1:]), walk=True, name="params to buffers")),
+  # LINEAR on LINEAR
+  (UPat(Ops.LINEAR, custom_early_reject={Ops.LINEAR}, name="x"),
+   lambda x: x.replace(src=tuple(flatten(x.src if x.op is Ops.LINEAR else (x,) for x in x.src)))),
 ])
 
 schedule_cache: dict[bytes, UOp] = {}
-def lower_schedule_to_linear(big_sink:UOp) -> UOp|None:
+# ctx is just for DEBUG on inner
+def lower_sink_to_linear(function:UOp) -> UOp|None:
   st = time.perf_counter()
-  function = big_sink.src[0]
   if isinstance(function.arg, KernelInfo): return None
-  if not SCACHE or (sc_ret:=schedule_cache.get(function.key, None)) is None:
-    if SPEC: type_verify(big_sink, tensor_spec)
+  cache_key = function.key
+  if not SCACHE or (sc_ret:=schedule_cache.get(cache_key, None)) is None:
+    if SPEC: type_verify(function, tensor_spec)
     # support recursive CALLs
-    function = graph_rewrite(function, pm_schedule, name="inner schedule to linear")
     linear = create_schedule(get_kernel_graph(function))
-    if SCACHE: schedule_cache[function.key] = linear
+    if SCACHE: schedule_cache[cache_key] = linear
   else:
     # schedule cache hit
     linear = sc_ret
@@ -124,20 +127,21 @@ def lower_schedule_to_linear(big_sink: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'} {function.key.hex()[:8]}"+\
+          f" | {' cache hit' if SCACHE and sc_ret is not None else 'CACHE MISS'} {cache_key.hex()[:8]}"+\
           f" | {len(UOpMetaClass.ucache):7d} uops in cache"+("" if frm is None else f" | {frm.filename}:{frm.lineno}"))
-  # 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")
+  return linear
 
 pm_schedule = PatternMatcher([
-  (UPat(Ops.CALL, src=(UPat(Ops.SINK),), allow_any_len=True, name="big_sink"), lower_schedule_to_linear),
+  (UPat(Ops.SINK, name="function"), lower_sink_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 = graph_rewrite(big_sink, pm_schedule, name="schedule to linear")
+  linear_call = graph_rewrite(big_sink, pm_schedule, name="schedule to linear", enter_calls=True)
+
+  # this recursively resolves the linear_call and allocates buffers
+  linear = graph_rewrite(linear_call, pm_resolve_linear_call, name="resolve linear call")
 
   # vars used in the schedule
   used_vars = set().union(*[{v.expr for v in si.src[0].variables()} for si in linear.src])
@@ -148,7 +152,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
-      assert nm not in var_vals or var_vals[nm] == val, f"bind mismatch on {nm}, {var_vals[nm]} != {val}"
+      if var_vals.get(nm, val) != val: raise RuntimeError(f"bind mismatch on {nm}, {var_vals[nm]} != {val}")
       var_vals[nm] = val
 
   # convert LINEAR to ExecItems

tinygrad/function.py

@@ -1,5 +1,5 @@
 import functools
-from typing import Generic, TypeVar, Callable, cast
+from typing import Generic, TypeVar, Callable, cast, overload
 from tinygrad.helpers import Context, dedup, getenv
 from tinygrad.uop.ops import UOp, Ops, graph_rewrite, PatternMatcher, UPat
 from tinygrad.tensor import Tensor
@@ -10,15 +10,16 @@ def add_to_ctx(ctx, x:UOp):
   return ret
 
 pm_ctx = PatternMatcher([
-  (UPat(Ops.BUFFER, name="x"), add_to_ctx),
+  (UPat((Ops.BUFFER, Ops.BIND), name="x"), add_to_ctx),
   (UPat((Ops.ASSIGN, Ops.CONTIGUOUS), name="x"),
    lambda ctx,x: add_to_ctx(ctx,x) if not x.op_in_backward_slice_with_self(Ops.PARAM) else None),
 ])
 
 ReturnType = TypeVar('ReturnType')
-class function(Generic[ReturnType]):
-  def __init__(self, fxn:Callable[..., ReturnType]):
+class _function(Generic[ReturnType]):
+  def __init__(self, fxn:Callable[..., ReturnType], *, precompile:bool=False):
     self.fxn = fxn
+    self.precompile = precompile
 
   def __get__(self, obj, objtype=None): return functools.partial(self.__call__, obj) if obj is not None else self
 
@@ -57,6 +58,14 @@ class function(Generic[ReturnType]):
     #call = assigned.call(*call_uops, buffer, name=name)
     #ret = buffer.after(call)
 
-    ret = uret.call(*call_uops, name=name)
+    ret = uret.call(*call_uops, name=name, precompile=self.precompile)
     return cast(ReturnType, Tensor(ret, device=ret.device))
 
+# overload signatures support both @function and @function(precompile=True) syntax
+@overload
+def function(fxn:Callable[..., ReturnType], *, precompile:bool=False) -> _function[ReturnType]: ...
+@overload
+def function(fxn:None=None, *, precompile:bool=False) -> Callable[[Callable[..., ReturnType]], _function[ReturnType]]: ...
+def function(fxn=None, *, precompile:bool=False):
+  if fxn is None: return lambda f: _function(f, precompile=precompile)
+  return _function(fxn, precompile=precompile)

tinygrad/gradient.py

@@ -17,12 +17,13 @@ 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 = 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))
+  params = {x.arg:x for x in fxn.toposort(enter_calls=False) if x.op == Ops.PARAM}
+  grads = compute_gradient(fxn, ctx.param_like(len(args)), set(params.values()))
   ret: list[UOp|None] = [None]
-  for i,p in enumerate(params):
-    if p in grads:
+  for i in range(len(args)):
+    if (p:=params.get(i, None)) is not None and p in grads:
       # TODO: compact the args and remove unused ones
+      assert not grads[p].op_in_backward_slice_with_self(Ops.BUFFER), "BUG: BUFFER in backward slice of grad"
       ret.append(grads[p].call(*args, ctx, name=(k.arg.name or "")+f"_backward_{i}"))
     else:
       ret.append(None)

tinygrad/helpers.py

@@ -1,5 +1,6 @@
 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
@@ -172,7 +173,8 @@ 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, IMAGE, BEAM, NOOPT = ContextVar("DEBUG", 0), ContextVar("IMAGE", 0), ContextVar("BEAM", 0), ContextVar("NOOPT", 0)
+DEBUG, BEAM, NOOPT = ContextVar("DEBUG", 0), ContextVar("BEAM", 0), ContextVar("NOOPT", 0)
+IMAGE, FLOAT16 = ContextVar("IMAGE", 0), ContextVar("FLOAT16", 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)
@@ -231,6 +233,7 @@ 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,6 +17,7 @@ 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,
@@ -24,10 +25,9 @@ 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]:
-    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]
+    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]
 
   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].bitcast(dtype)
+    return t[:dtype.itemsize * n].contiguous().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]))
+    blocks = t[:(n//nelements_nbytes[0])*nelements_nbytes[1]].reshape((-1, nelements_nbytes[1])).contiguous()
     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,6 +46,7 @@ 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
@@ -72,8 +73,10 @@ 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)
@@ -99,19 +102,36 @@ 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
-  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
+  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
 
 # ═══════════════════════════════════════════════════════════════════════════════
 # PACKET TYPE BASE CLASS
@@ -147,11 +167,6 @@ 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]
@@ -341,14 +356,9 @@ class INST(PacketType):
 class INST_RDNA4(PacketType):  # Layout 4: different delta position and InstOp encoding
   encoding = bits[2:0] == 0b010
   delta = bits[5:3]
-  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
+  w64h = bits[6:6]
+  wave = bits[11:7]
+  op = bits[19:12].enum(InstOpRDNA4)
 
 class UTILCTR(PacketType):
   encoding = bits[6:0] == 0b0110001
@@ -363,7 +373,7 @@ PACKET_TYPES_RDNA3: dict[int, type[PacketType]] = {
 }
 PACKET_TYPES_RDNA4: dict[int, type[PacketType]] = {
   **PACKET_TYPES_RDNA3,
-  7: TS_DELTA_S8_W3_RDNA4, 9: WAVESTART_RDNA4, 10: TS_DELTA_S5_W2_RDNA4, 11: WAVEALLOC_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,
 }
 
@@ -536,8 +546,9 @@ 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):
-      chunk = int.from_bytes(data[pos:pos + byte_count], 'little')
-      reg, pos = (reg >> (byte_count * 8)) | (chunk << (64 - byte_count * 8)), pos + byte_count
+      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
     # 3. if odd, read low nibble
     if (nib_off := need & 1): reg = (reg >> 4) | ((data[pos] & 0xF) << 60)
 
@@ -619,9 +630,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", "NEXT"}, f"branch can only be folowed by JUMP, got {p}"
+        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}"
       # JUMP handling
-      if (isinstance(p, INST) and p.op is InstOp.JUMP) or (isinstance(p, INST_RDNA4) and branch_inst is not None and p.flag3):
+      if (isinstance(p, INST) and p.op is InstOp.JUMP) or (isinstance(p, INST_RDNA4) and p.op is InstOpRDNA4.JUMP):
         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
@@ -650,7 +661,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 "")
+    fields = f"wave={p.wave} op={op_name}" + ((" flag1" if p.flag1 else "") + (" flag2" if p.flag2 else "") if isinstance(p, INST) 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}"
@@ -666,18 +677,19 @@ 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 p in packets:
-    if type(p).__name__.replace("_RDNA4", "") not in skip: print(format_packet(p))
+  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 '')
 
 if __name__ == "__main__":
   import sys, pickle
-  if len(sys.argv) < 2:
-    print("Usage: python sqtt.py <pkl_file>")
-    sys.exit(1)
-  with open(sys.argv[1], "rb") as f:
+  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:
     data = pickle.load(f)
-  prg_names = {e.tag: e.name for e in data if type(e).__name__ == "ProfileProgramEvent" and e.tag is not None}
+  prg_events = {e.tag: e 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):
-    print(f"\n=== event {i} {prg_names.get(event.kern, '')} ===")
-    print_packets(decode(event.blob))
+    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))

tinygrad/runtime/ops_amd.py

@@ -598,9 +598,10 @@ 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):
+  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):
     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)
+    res = super().__call__(*bufs, global_size=global_size, local_size=local_size, vals=vals, wait=wait, timeout=timeout)
     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)
@@ -857,21 +858,21 @@ class PCIIface(PCIIfaceBase):
 
     rcvr_params: tuple
     if queue_type == kfd.KFD_IOC_QUEUE_TYPE_SDMA:
-      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)))
+      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:
-      pv, doorbell_index = self.dev_impl.gfx.setup_ring(*(rcvr_params:=(ring.va_addr, ring.size, gart.va_addr+rptr, gart.va_addr+wptr,
+      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=pv,
+    return AMDQueueDesc(ring=ring.cpu_view().view(fmt='I'), doorbell=self.dev_impl.doorbell64.view(doorbell_index * 8, 8, fmt='Q'), put_value=0,
       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():
-        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
+      if reset and d.iface.dev_impl.recover(force=d.error_state is not None):
+        d.compute_queue.put_value = d.compute_queue.read_ptr[0] = d.compute_queue.write_ptr[0] = 0
+        d.iface.dev_impl.gfx.setup_ring(*d.compute_queue.params)
         d.timeline_signal.value = d.timeline_value - 1
         d.error_state = None
 
@@ -977,7 +978,8 @@ 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)
+                     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())
 
     # 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) -> float|None:
+               vals:tuple[int, ...]=(), wait=False, **kw) -> 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):
+  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):
     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):
+  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):
     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):
+  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):
     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):
+  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):
     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):
+  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):
     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):
+  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):
     with cpu_profile(self.name, self.device): return 1e-3
 
 class NullAllocator(Allocator['NullDevice']):

tinygrad/runtime/ops_nv.py

@@ -312,12 +312,13 @@ 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):
+  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):
     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)
+    res = super().__call__(*bufs, global_size=global_size, local_size=local_size, vals=vals, wait=wait, timeout=timeout)
     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):
+  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):
     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,7 +266,8 @@ 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):
+  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):
     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) -> float|None:
+               vals:tuple[int, ...]=(), wait=False, **kw) -> 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, self.has_aql_queue = smi_dev, False, False
+    self.smi_dev, self.is_err_state = smi_dev, 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) -> 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")
+  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")
     self.ih.interrupt_handler()
     self.gfx.reset_mec()
     self.is_err_state = False
-    if DEBUG >= 2: print(f"am {self.devfmt}: Recovery complete")
+    if DEBUG >= 3: 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_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}}
+    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}}
 
   def init_hw(self):
     if self.adev.ip_ver[am.NBIO_HWIP] in {(7,9,0), (7,9,1)}:
@@ -240,10 +240,11 @@ class AM_GFX(AM_IP):
 
   def init_hw(self):
     # Wait for RLC autoload to complete
-    while self.adev.regCP_STAT.read() != 0 and self.adev.regRLC_RLCS_BOOTLOAD_STATUS.read_bitfields()['bootload_complete'] != 0: pass
+    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")
 
     self.adev.gmc.init_hub("GC", inst_cnt=self.xccs)
-    if self.adev.partial_boot: return
+    if self.adev.partial_boot: return self.reset_mec()
 
     self._config_mec()
 
@@ -285,24 +286,26 @@ class AM_GFX(AM_IP):
     self._enable_mec()
 
     # Set 1 partition
-    if self.xccs > 1 and not self.adev.partial_boot: self.adev.psp._spatial_partition_cmd(1)
+    if self.xccs > 1: self.adev.psp._spatial_partition_cmd(1)
 
   def fini_hw(self): self._dequeue_hqds()
 
   def reset_mec(self):
-    self._dequeue_hqds(reset=True)
+    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._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) -> tuple[int, int]:
-    self.adev.has_aql_queue |= aql
+  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:
     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,
@@ -320,26 +323,16 @@ 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)
 
-      # Copy mqd into memory
-      self._grbm_select(me=1, pipe=pipe, queue=queue, inst=xcc)
+      self.adev.vram.view(self.mqd_paddr[queue] + 0x1000*xcc, ctypes.sizeof(mqd_struct))[:] = memoryview(mqd_struct).cast('B')
 
-      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)
+      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 restore_ptr // 16, doorbell
+    return 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)
@@ -390,15 +383,13 @@ 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, 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):
+  def _dequeue_hqds(self):
+    for q in range(2):
       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 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")
+          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")
     self._grbm_select()
 
 class AM_IH(AM_IP):
@@ -445,7 +436,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_scrs_names.get(client, {}).get(src, '')
+      src_name = self.adev.soc.ih_srcs_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}]")
 
@@ -459,6 +450,7 @@ 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
 
@@ -515,21 +507,23 @@ 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) -> tuple[int, int]:
+  def setup_ring(self, ring_addr:int, ring_size:int, rptr_addr:int, wptr_addr:int, idx: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)
-    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_RPTR", "", "_HI", 0, inst=inst)
+    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)
@@ -539,7 +533,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 self.adev.reg(f"{reg}_RB_WPTR").read(inst=inst) | (self.adev.reg(f"{reg}_RB_WPTR_HI").read(inst=inst) << 32), doorbell
+    return 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=getenv("HCQDEV_WAIT_TIMEOUT_MS", 30000)):
+  def wait(self, value:int, timeout:int|None=None):
     """
     Waits the signal is greater than or equal to a specific value.
 
@@ -261,6 +261,7 @@ 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)
@@ -325,7 +326,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) -> float|None:
+               vals:tuple[int|None, ...]=(), wait:bool=False, timeout:int|None=None) -> float|None:
     """
     Enqueues the program for execution with the given arguments and dimensions.
 
@@ -349,7 +350,7 @@ class HCQProgram(Generic[HCQDeviceType]):
 
     q.signal(self.dev.timeline_signal, self.dev.next_timeline()).submit(self.dev)
 
-    if wait: self.dev.synchronize()
+    if wait: self.dev.synchronize(timeout=timeout)
     return (float(sig_en.timestamp - sig_st.timestamp) / 1e6) if wait else None
 
 class HCQCompiled(Compiled, Generic[SignalType]):
@@ -362,7 +363,8 @@ 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):
+               comp_queue_t:Callable[..., HWQueue], copy_queue_t:Callable[..., HWQueue]|None=None, kernargs_size=(16 << 20), sigalloc_size=0x1000,
+               can_recover:bool=False):
     self.device_id:int = int(device.split(":")[1]) if ":" in device else 0
 
     from tinygrad.runtime.graph.hcq import HCQGraph
@@ -386,22 +388,23 @@ 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):
+  def synchronize(self, timeout:int|None=None):
     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)
+    try: self.timeline_signal.wait(self.timeline_value - 1, timeout=timeout if timeout is not None and self.can_recover else None)
     except RuntimeError as e:
       self.error_state = e
       if hasattr(self, 'on_device_hang'): self.on_device_hang()
-      else: raise e
+      raise e
 
     if self.timeline_value > (1 << 31): self._wrap_timeline_signal()
     if PROFILE:

tinygrad/schedule/allreduce.py

@@ -0,0 +1,63 @@
+import functools, itertools
+from tinygrad.helpers import all_int, prod, DEBUG, RING, ALL2ALL, getenv
+from tinygrad.uop.ops import Ops, UOp
+
+# *** allreduce implementation ***
+def handle_allreduce(buf:UOp, red:UOp) -> UOp|None:
+  if not isinstance(buf.device, tuple): return None
+  assert all_int(buf.shape), f"does not support symbolic shape {buf.shape}"
+  ndev, shape, numel = len(buf.device), buf.shape, prod(buf.shape)
+
+  # ring allreduce doesn't provide a benefit with only 2 nodes or where number of elements is less than 256k (empirically)
+  # fallback to naive allreduce to save on kernel dispatch, chunking and reassembling chunks.
+  use_all2all = (ALL2ALL >= 2 or (ndev > 2 and numel > getenv("RING_ALLREDUCE_THRESHOLD", 256_000) and ALL2ALL >= 1))
+  use_ring = not use_all2all and (RING >= 2 or (ndev > 2 and numel > getenv("RING_ALLREDUCE_THRESHOLD", 256_000) and RING >= 1))
+  if DEBUG >= 2: print(f"{'ALL2ALL' if use_all2all else 'RING' if use_ring else 'NAIVE'} ALLREDUCE {ndev}x{numel} | {buf.dtype}")
+
+  # contiguous before we copy it
+  buf = buf.contiguous()
+
+  # naive: copy to all devices. if you shrink later, that'll be handled
+  if not use_ring and not use_all2all:
+    return functools.reduce(lambda x,y: x.alu(red.arg, y), [UOp(Ops.COPY, buf.dtype, (buf.mselect(i), red.src[1])) for i in range(ndev)])
+
+  # chunk data into ndev pieces
+  factor = next((f for f in [32, 16, 8, 4, 2] if numel % f == 0), 1)
+  base, left = divmod(numel // factor,  ndev)
+  chunks = list(itertools.pairwise(itertools.accumulate([(base + 1) * factor] * left + [base * factor] * (ndev - left), initial=0)))
+
+  # reduce-scatter
+  reduced_chunks:list[UOp] = []
+  for i,(s,e) in enumerate(chunks):
+    if use_all2all:
+      chunks_on_i = [buf.mselect(j).reshape((numel,)).shrink(((s,e),)).copy_to_device(buf.device[i]) for j in range(ndev)]
+      reduced_chunks.append(functools.reduce(lambda x,y: x.alu(red.arg, y), chunks_on_i))
+    else:
+      chunk, reduced = buf.reshape((numel,)).shrink(((s,e),)), buf.reshape((numel,)).shrink(((s,e),))
+      for step in range(ndev-1):
+        src, dest = (i+step)%ndev, (i+step+1)%ndev
+        cp = reduced.copy_to_device(buf.device[dest], src if isinstance(reduced.device, tuple) else None)
+        reduced = cp.alu(red.arg, chunk.copy_to_device(buf.device[dest], dest))
+      reduced_chunks.append(reduced)
+
+  # allgather
+  copied_chunks:list[UOp] = []
+  for i,rc in enumerate(reduced_chunks):
+    if isinstance(red.src[1].arg, str): copied_chunks.append(rc.copy_to_device(red.src[1].arg))
+    elif use_all2all: copied_chunks.append(UOp(Ops.MSTACK, buf.dtype, tuple(rc.copy_to_device(buf.device[j]) for j in range(ndev))))
+    else:
+      chain:list[UOp] = [rc]
+      for step in range(ndev-1):
+        chain.append(rc := rc.copy_to_device(buf.device[(i+step)%ndev]))
+      copied_chunks.append(UOp(Ops.MSTACK, buf.dtype, tuple(chain[(j-i+1)%ndev] for j in range(ndev))))
+
+  # reassemble
+  return UOp.sum(*[c.pad(((s,numel-e),)) for (s,e),c in zip(chunks, copied_chunks)]).reshape(shape)
+
+def create_allreduce_function(buf:UOp, red:UOp, output:UOp|None=None) -> UOp|None:
+  # BUFFER without unique have unique added later
+  if output is None: output = UOp(Ops.BUFFER, red.dtype, (UOp(Ops.NOOP), red.src[1]), red.size).reshape(red.shape)
+  to = red.param_like(0)
+  src = buf.param_like(1)
+  red = UOp(Ops.ALLREDUCE, dtype=red.dtype, src=(src, red.src[1]), arg=red.arg)
+  return output.after(to.assign(handle_allreduce(src, red)).sink().call(output, buf.contiguous(), name="allreduce", precompile=True))

tinygrad/schedule/indexing.py

@@ -56,7 +56,7 @@ class IndexingContext:
     return UOp.range(s, next(self.range_idx), axistype) if resolve(s!=1) else UOp.const(dtypes.index, 0)
 
 def create_bufferize_and_index_based_on_ranges(ctx:IndexingContext, x:UOp):
-  if x.op in {Ops.BUFFERIZE, Ops.INDEX, Ops.AFTER}: return None
+  if x.op in {Ops.BUFFERIZE, Ops.INDEX}: return None
   new_srcs = []
   for s in x.src:
     new_src = s
@@ -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)
+      sink = UOp.sink(*rngs).simplify() # NOTE: this applies any commutative flips to the rngs early
       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")
@@ -179,6 +179,9 @@ def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]:
     # no ranges on kernels, they are internal
     if x.op in {Ops.CALL, Ops.LINEAR}: continue
 
+    # no range on after
+    if x.op is Ops.AFTER: continue
+
     # treat MSTACK/MSELECT like SINK
     if x.op in {Ops.MSTACK, Ops.MSELECT}: continue
 

tinygrad/schedule/multi.py

@@ -1,59 +1,7 @@
-import functools, itertools
-from tinygrad.helpers import all_same, all_int, prod, DEBUG, RING, ALL2ALL, getenv
+from tinygrad.helpers import all_same, prod, getenv
 from tinygrad.uop.ops import Ops, UOp, PatternMatcher, UPat, GroupOp, graph_rewrite, should_resolve_call
 from tinygrad.dtype import dtypes
-
-# *** allreduce implementation ***
-def handle_allreduce(buf:UOp, red:UOp) -> UOp|None:
-  if not isinstance(buf.device, tuple): return None
-  assert all_int(buf.shape), f"does not support symbolic shape {buf.shape}"
-  ndev, shape, numel = len(buf.device), buf.shape, prod(buf.shape)
-
-  # ring allreduce doesn't provide a benefit with only 2 nodes or where number of elements is less than 256k (empirically)
-  # fallback to naive allreduce to save on kernel dispatch, chunking and reassembling chunks.
-  use_all2all = (ALL2ALL >= 2 or (ndev > 2 and numel > getenv("RING_ALLREDUCE_THRESHOLD", 256_000) and ALL2ALL >= 1))
-  use_ring = not use_all2all and (RING >= 2 or (ndev > 2 and numel > getenv("RING_ALLREDUCE_THRESHOLD", 256_000) and RING >= 1))
-  if DEBUG >= 2: print(f"{'ALL2ALL' if use_all2all else 'RING' if use_ring else 'NAIVE'} ALLREDUCE {ndev}x{numel} | {buf.dtype}")
-
-  # contiguous before we copy it
-  buf = buf.contiguous()
-
-  # naive: copy to all devices. if you shrink later, that'll be handled
-  if not use_ring and not use_all2all:
-    return functools.reduce(lambda x,y: x.alu(red.arg, y), [UOp(Ops.COPY, buf.dtype, (buf.mselect(i), red.src[1])) for i in range(ndev)])
-
-  # chunk data into ndev pieces
-  factor = next((f for f in [32, 16, 8, 4, 2] if numel % f == 0), 1)
-  base, left = divmod(numel // factor,  ndev)
-  chunks = list(itertools.pairwise(itertools.accumulate([(base + 1) * factor] * left + [base * factor] * (ndev - left), initial=0)))
-
-  # reduce-scatter
-  reduced_chunks:list[UOp] = []
-  for i,(s,e) in enumerate(chunks):
-    if use_all2all:
-      chunks_on_i = [buf.mselect(j).reshape((numel,)).shrink(((s,e),)).copy_to_device(buf.device[i]) for j in range(ndev)]
-      reduced_chunks.append(functools.reduce(lambda x,y: x.alu(red.arg, y), chunks_on_i))
-    else:
-      chunk, reduced = buf.reshape((numel,)).shrink(((s,e),)), buf.reshape((numel,)).shrink(((s,e),))
-      for step in range(ndev-1):
-        src, dest = (i+step)%ndev, (i+step+1)%ndev
-        cp = reduced.copy_to_device(buf.device[dest], src if isinstance(reduced.device, tuple) else None)
-        reduced = cp.alu(red.arg, chunk.copy_to_device(buf.device[dest], dest))
-      reduced_chunks.append(reduced)
-
-  # allgather
-  copied_chunks:list[UOp] = []
-  for i,rc in enumerate(reduced_chunks):
-    if isinstance(red.src[1].arg, str): copied_chunks.append(rc.copy_to_device(red.src[1].arg))
-    elif use_all2all: copied_chunks.append(UOp(Ops.MSTACK, buf.dtype, tuple(rc.copy_to_device(buf.device[j]) for j in range(ndev))))
-    else:
-      chain:list[UOp] = [rc]
-      for step in range(ndev-1):
-        chain.append(rc := rc.copy_to_device(buf.device[(i+step)%ndev]))
-      copied_chunks.append(UOp(Ops.MSTACK, buf.dtype, tuple(chain[(j-i+1)%ndev] for j in range(ndev))))
-
-  # reassemble
-  return UOp.sum(*[c.pad(((s,numel-e),)) for (s,e),c in zip(chunks, copied_chunks)]).reshape(shape)
+from tinygrad.schedule.allreduce import handle_allreduce
 
 # ***** multi rewrite MSELECT/MSTACK *****
 
@@ -71,7 +19,6 @@ 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),
   # 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),
@@ -87,6 +34,11 @@ replace_allreduce = PatternMatcher([
    lambda s,v,ms: v.replace(src=(s.mselect(ms.arg),)+v.src[1:])),
 ])
 
+_early_allreduce = PatternMatcher([
+  (UPat(Ops.ALLREDUCE, src=(UPat.var("buf"), UPat()), name="red"), handle_allreduce),
+])
+if not getenv("LATE_ALLREDUCE", 0): replace_allreduce = _early_allreduce + replace_allreduce
+
 # ***** multi functions *****
 
 def alu_multi(root:UOp):
@@ -164,10 +116,18 @@ 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):
-  return call.replace(src=(graph_rewrite(call.src[0], multi_pm, name="subcall"),)+call.src[1:]) if should_resolve_call(call) else None
+  if not should_resolve_call(call): return None
+  new_body = graph_rewrite(call.src[0], multi_pm, name="subcall")
+  new_args = tuple(a.src[0] if a.op is Ops.MULTI else a for a in call.src[1:])
+  return call.replace(src=(new_body,)+new_args)
+
+def param_to_multi(p:UOp):
+  if p.axis is None: return None
+  return UOp.param(p.arg, p.dtype, p.shard_shape, p._device).multi(p.axis)
 
 # NOTE: this is the same pattern as Ops.UNROLL
 multi_pm = PatternMatcher([
+  (UPat(Ops.PARAM, name="p"), param_to_multi),
   (UPat(GroupOp.ALU, name="root", custom_early_reject=set([Ops.MULTI])), alu_multi),
   (UPat(Ops.REDUCE_AXIS, src=(UPat(Ops.MULTI, name="multi"), ), name="root"), reduce_multi),
   (UPat(Ops.RESHAPE, src=(UPat(Ops.MULTI, name="multi"), UPat()), name="root"), reshape_multi),

tinygrad/schedule/rangeify.py

@@ -1,15 +1,16 @@
 from dataclasses import dataclass, field, replace
 import itertools
-from tinygrad.dtype import dtypes, PtrDType, ImageDType, AddrSpace
+from tinygrad.dtype import dtypes, PtrDType, ImageDType, AddrSpace, Invalid
 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
+from tinygrad.uop.ops import graph_rewrite, sint, AxisType, BottomUpGate, profile_matches, should_resolve_call, identity_element
 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
+from tinygrad.helpers import PCONTIG, FLOAT16, argsort, partition, get_single_element
 from tinygrad.codegen.simplify import pm_flatten_range, pm_reduce_simplify
 from tinygrad.codegen.opt import Opt
 from tinygrad.schedule.indexing import run_rangeify, BufferizeOpts, ALWAYS_CONTIGUOUS, IndexingContext, apply_movement_op
 from tinygrad.schedule.multi import multi_pm
+from tinygrad.schedule.allreduce import create_allreduce_function
 
 # creation can recurse a lot
 import sys
@@ -21,6 +22,22 @@ pm_syntactic_sugar = PatternMatcher([
    lambda i1,i2: i2.replace(src=i1.src+i2.src[1:]) if isinstance(i1.dtype, PtrDType) and not isinstance(i2.dtype, PtrDType) else None),
 ])
 
+def found_assign(ctx:dict[UOp, UOp], assign:UOp, src:UOp):
+  if (x:=src).op is Ops.CAST and x.dtype == dtypes.half and FLOAT16: x, assign = x.src[0], assign.cast(dtypes.float)
+  while x is not x.base:
+    if x.op is Ops.PERMUTE: assign = assign.permute(argsort(x.marg))
+    elif x.op is Ops.RESHAPE: assign = assign.reshape(x.src[0].shape)
+    else: return None
+    x = x.src[0]
+  ctx[x] = assign
+
+pm_openpilot_hacks = PatternMatcher([
+  # *** ASSIGN replacement hack for openpilot ***
+  (UPat(Ops.ASSIGN, src=(UPat(), UPat((*GroupOp.Movement, Ops.CAST), name="src")), name="assign"), found_assign),
+  # replace ALU sources with assign 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),
+])
+
 # movement op on INDEX as a PatternMatcher
 pm_mops = PatternMatcher([
   (UPat(GroupOp.Movement, name="r").f(Ops.INDEX, allow_any_len=True, name="idx"),
@@ -90,14 +107,23 @@ 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.shape != a.shape: raise TypeError(f"arg {i} shape mismatch: expected {p.shape}, got {a.shape}")
+    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),
 
+  # resolve allreduce (must be bottom up)
+  (UPat(Ops.ASSIGN, src=(UPat.var("output"), UPat(Ops.ALLREDUCE, src=(UPat.var("buf"), UPat()), name="red"))), create_allreduce_function),
+  (UPat(Ops.ALLREDUCE, src=(UPat.var("buf"), UPat()), name="red"), create_allreduce_function),
+
   # split_reduceop
   (UPat(Ops.REDUCE_AXIS, name="reduce", src=(UPat.var("x"),)), split_reduceop),
 
@@ -135,6 +161,14 @@ 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),
 ])
 
 # *****************
@@ -229,8 +263,9 @@ 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
-  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)
+  # 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)
 
 def remove_noop_bufferize(idx,b2):
   if idx.src[1:] != b2.src[1:] or idx.src[0].op is Ops.BUFFER_VIEW: return None
@@ -355,6 +390,12 @@ def flatten_bufferize(x:UOp):
   return ret
 pm_flatten_bufferize = PatternMatcher([(UPat(Ops.BUFFERIZE, name="x"), flatten_bufferize)])
 
+def resolve_anonymous_buffer(ctx:itertools.count, b:UOp, c:UOp) -> UOp|None:
+  dab = b.replace(src=(UOp(Ops.LUNIQUE, arg=next(ctx)),)+b.src[1:])
+  nc_src = tuple(dab if x is b else x for x in c.src)
+  if nc_src == c.src: return None
+  return dab.after(c.replace(src=nc_src))
+
 pm_add_buffers = pm_mops+pm_flatten_bufferize+to_bufferview+PatternMatcher([
   (UPat(Ops.BUFFERIZE, src=(UPat(), UPat(name="idx")), name="x"), lambda ctx,x,idx: bufferize_to_store(ctx, x, idx, allow_locals=False)),
 
@@ -364,6 +405,14 @@ pm_add_buffers = pm_mops+pm_flatten_bufferize+to_bufferview+PatternMatcher([
 
   # remove any RESHAPEs on KERNEL
   (UPat(Ops.CALL, name="k"), lambda k: k.replace(src=tuple(x.src[0] if x.op is Ops.RESHAPE else x for x in k.src))),
+
+  # remove MOP on AFTER
+  (UPat(Ops.AFTER, src=(UPat.var("x"), UPat(GroupOp.Movement, name="y"))), lambda x,y: x.after(y.src[0])),
+  # remove double AFTER
+  (UPat(Ops.AFTER, src=(UPat.var("x"), UPat(Ops.AFTER, name="y"))), lambda x,y: x.after(*y.src[1:])),
+
+  # resolve anonymous buffers
+  (UPat(Ops.AFTER, src=(UPat(Ops.BUFFER, src=(UPat(Ops.NOOP),), name="b", allow_any_len=True), UPat(Ops.CALL, name="c"))), resolve_anonymous_buffer),
 ])
 
 pm_add_buffers_local = pm_mops+pm_flatten_bufferize+to_bufferview+PatternMatcher([
@@ -382,7 +431,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)
+  ret = UOp(Ops.PARAM, buf.dtype.ptr(buf.size), arg=ctx.dg).reshape(buf.shape)
   if buf not in ctx.map: ctx.map[buf] = buf
   ctx.dg += 1
   return ret
@@ -486,7 +535,8 @@ split_kernels = PatternMatcher([
 
 @profile_matches
 def get_kernel_graph(sink:UOp) -> UOp:
-  tsink = graph_rewrite(sink, multi_pm, name="multi_pm")
+  tsink = graph_rewrite(sink, pm_openpilot_hacks, ctx={}, name="openpilot hacks")
+  tsink = graph_rewrite(tsink, multi_pm, name="multi_pm")
   tsink = graph_rewrite(tsink, pm_syntactic_sugar+pm_mops+earliest_rewrites, bottom_up=True, name="earliest rewrites")
 
   # convert movement ops to ranges

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, WINO, Metadata, TRACEMETA, ASM_GEMM, ceildiv, fetch, is_numpy_ndarray, TracingKey, cpu_profile
+from tinygrad.helpers import IMAGE, FLOAT16, 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,8 +25,7 @@ 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] = {}
-_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:
+def _apply_map_to_tensors(applied_map:dict[UOp, UOp], name:str, walk:bool=False) -> None:
   with cpu_profile(TracingKey(name), "TINY"):
     # get tensors in scope
     in_scope: dict[UOp, bool] = {}
@@ -35,7 +34,7 @@ def _apply_map_to_tensors(applied_map:dict[UOp, UOp], name:str) -> None:
 
     # 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}")
+    new_sink = sink.substitute(applied_map, name=f"substitute {name}", walk=walk)
 
     # set the relevant uop to the realized UOps
     for t,s,ns in zip(scope_tensors, sink.src, new_sink.src):
@@ -278,23 +277,6 @@ 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
@@ -323,13 +305,14 @@ class Tensor(OpMixin):
     if is_disk:
       self._buffer().copyin(x._data())
       return self
-    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)
+    # 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 not in {Ops.BUFFER, Ops.AFTER, Ops.ASSIGN}: raise RuntimeError(f"can only assign to buffer, not {base.op}")
+    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))
 
   def detach(self) -> Tensor:
     """
@@ -639,7 +622,7 @@ class Tensor(OpMixin):
       Tensor._device_seeds[device] = Tensor(
         [int.from_bytes(hashlib.sha256(len(Tensor._device_seeds).to_bytes(4, "big")).digest(), "big"), Tensor._seed],
         device=device, dtype=dtypes.uint32, requires_grad=False)
-      Tensor._device_rng_counters[device] = Tensor([num], device=device, dtype=dtypes.uint32, requires_grad=False).contiguous()
+      Tensor._device_rng_counters[device] = Tensor([num], device=device, dtype=dtypes.uint32, requires_grad=False).clone()
     # increment rng counter for devices
     else: Tensor._device_rng_counters[device].assign(Tensor._device_rng_counters[device] + num)
 
@@ -1236,26 +1219,6 @@ 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]), []
@@ -1266,7 +1229,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 e: raise IndexError(f"cannot broadcast indices: {e}") from e
+        except ValueError as err: raise IndexError(f"cannot broadcast indices: {err}") from err
         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),)]
@@ -1277,7 +1240,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 e: raise IndexError(f"cannot broadcast indices: {e}") from e
+        except ValueError as err: raise IndexError(f"cannot broadcast indices: {err}") from err
         masks.append(i._one_hot_along_dim(num_classes=x.shape[dim], dim=(dim - x.ndim)))
 
       # reduce masks to 1 mask
@@ -1286,21 +1249,36 @@ 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))
 
-      # 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
+      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)
 
   def __getitem__(self, indices) -> Tensor:
     """
@@ -1344,15 +1322,26 @@ 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): raise NotImplementedError("setitem with requires_grad is not supported")
+    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
     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: # basic setitem, self is realized. TODO: disk uop.base is a COPY and not realized
-      self[indices].assign(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)
     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
@@ -1762,7 +1751,7 @@ class Tensor(OpMixin):
     print(t.all(axis=1, keepdim=True).numpy())
     ```
     """
-    return self.logical_not().any(axis, keepdim).logical_not()
+    return self.bool().min(axis, keepdim)
 
   def isclose(self, other:Tensor, rtol:float=1e-05, atol:float=1e-08, equal_nan=False) -> Tensor:
     """
@@ -2577,11 +2566,10 @@ class Tensor(OpMixin):
     return values._inverse(), indices
 
   @staticmethod
-  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=}"
+  def _tri(r:sint, c:sint, diagonal=0, device=None, requires_grad:bool|None=None) -> Tensor:
     return (Tensor.arange(r, device=device).unsqueeze(-1) + diagonal <= Tensor.arange(c, device=device)).requires_grad_(requires_grad)
 
-  def triu(self, diagonal:int=0) -> Tensor:
+  def triu(self, diagonal:sint=0) -> Tensor:
     """
     Returns the upper triangular part of the tensor, the other elements are set to 0.
 
@@ -2604,7 +2592,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:int=0) -> Tensor:
+  def tril(self, diagonal:sint=0) -> Tensor:
     """
     Returns the lower triangular part of the tensor, the other elements are set to 0.
 
@@ -3278,7 +3266,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()+1).item())
+    if num_classes == -1: num_classes = int(self.max().item())+1
     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,
@@ -3296,9 +3284,6 @@ 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)
@@ -3474,8 +3459,9 @@ 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)
-    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))
+    # 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()
     #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)
@@ -3613,7 +3599,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:=getenv("FLOAT16", 0)) else (dtypes.imagef, 4)
+    base_image_type, dtsz = (dtypes.imageh, 2) if FLOAT16 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

@@ -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) -> dict[UOp, None]:
+  def toposort(self, gate:Callable|None=None, enter_calls=True) -> dict[UOp, None]:
     cache: dict[UOp, None] = {}
     stack: list[tuple[UOp, bool]] = [(self, False)] # each stack entry is (node, visited_flag)
     while stack:
@@ -172,7 +172,8 @@ 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): stack.append((s, False)) # push srcs on the stack
+          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
       else: cache[node] = None # second time i'm seeing this node, add it to returned toposort
     return cache
 
@@ -253,6 +254,9 @@ 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}):
@@ -514,6 +518,11 @@ 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
@@ -648,34 +657,45 @@ 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.PARAM}
+    return self.op in {Ops.BUFFER, Ops.BUFFER_VIEW, Ops.PARAM}
 
   @property
   def buffer(self) -> Buffer|MultiBuffer:
     from tinygrad.device import Buffer, MultiBuffer
-    if self.op in {Ops.CONTIGUOUS, Ops.RESHAPE}: return self.src[0].buffer
+    if self.op in {Ops.CONTIGUOUS, Ops.RESHAPE, Ops.DETACH, Ops.AFTER}: return self.src[0].buffer
     # this buffer can process disk tensors and simple movement ops
     if self is not self.base:
-      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
+      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
       assert isinstance(buf, Buffer), "must be a Buffer for movement ops"
-      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")
+      return buf.view(self.size, self.dtype, offset*self.dtype.itemsize)
     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)
@@ -867,12 +887,14 @@ 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])
-    return UOp.param(slot, self.dtype, self._shape, self._device)
+    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
 
-  def call(self, *srcs:UOp, grad_fxn:Callable|None=None, metadata:tuple[Metadata, ...]=(), name:str|None=None) -> UOp:
+  def call(self, *srcs:UOp, grad_fxn:Callable|None=None, metadata:tuple[Metadata, ...]=(), name:str|None=None, precompile:bool=False) -> 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))
+    return UOp(Ops.CALL, self.dtype, (self,)+srcs, CallInfo(grad_fxn, metadata, name, precompile))
   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)]
@@ -895,15 +917,16 @@ 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))
-  def __repr__(self): return f"CallInfo({id(self.grad_fxn) if self.grad_fxn else None}, {self.metadata}, {repr(self.name)})"
+  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 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 is Ops.PROGRAM: return False
-  if c.src[0].op is Ops.COPY: return False
+  if c.src[0].op in {Ops.PROGRAM, Ops.LINEAR, Ops.COPY}: return False
+  if c.arg.precompile: return False
   return True
 
 # ******** ops in python ********
@@ -1250,12 +1273,13 @@ 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):
+  def __init__(self, pm, bpm, ctx=None, enter_calls=False):
     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)
@@ -1278,7 +1302,7 @@ class RewriteContext:
           continue
         # no rewrite, process children then come back to rebuild
         stack.append((n, True))
-        if n.op is Ops.CALL: self.replace[n.src[0]] = n.src[0]
+        if not self.enter_calls and 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:
@@ -1318,7 +1342,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 new_n.op is Ops.CALL: self.replace[new_n.src[0]] = new_n.src[0]
+        if not self.enter_calls and 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))
@@ -1357,8 +1381,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) -> UOp:
-  rewrite_ctx = RewriteContext(pm if not bottom_up else None, pm if bottom_up else bpm, ctx)
+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)
   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)
@@ -1392,7 +1416,10 @@ 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: return not (x.op is Ops.SINK and isinstance(x.arg, KernelInfo))
+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 do_unbind(ctx:dict[Variable, int], x:UOp):
   v,i = x.unbind()
@@ -1422,6 +1449,7 @@ 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)),
@@ -1483,7 +1511,8 @@ 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),
-  (UPat(GroupOp.Movement, name="x"), lambda ctx,x: f"{ctx[x.src[0]]}.{x.op.name.lower()}({render_marg(ctx,x)})"),
+  # 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)})"),
   # 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"),

tinygrad/uop/spec.py

@@ -87,6 +87,9 @@ _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),
 

tinygrad/uop/symbolic.py

@@ -258,7 +258,8 @@ symbolic = symbolic_simple+commutative+PatternMatcher([
   ((UPat.var("x", dtypes.index) + UPat.cvar("c")).cast(dtypes.sints, name="cast"), lambda x,c,cast:x.cast(cast.dtype)+c.cast(cast.dtype)),
   # only RANGE/IF/STORE/KERNEL have side effects
   (UPat(Ops.AFTER, name="x"), lambda x: x.replace(src=(x.src[0],)+
-    tuple(flatten([(y,) if y.op in {Ops.RANGE, Ops.STORE, Ops.CALL, Ops.BARRIER, Ops.END, Ops.UNROLL, Ops.LINEAR} else y.src for y in x.src[1:]])))),
+    tuple(flatten([(y,) if y.op in {Ops.RANGE, Ops.STORE, Ops.CALL, Ops.BARRIER, Ops.END, Ops.UNROLL, Ops.LINEAR, Ops.BUFFERIZE}
+                   else y.src for y in x.src[1:]])))),
   # after with 1 src is just src[0]
   (UPat(Ops.AFTER, src=(UPat.var("s"),)), lambda s: s),
   # VECTORIZE/CONST
@@ -426,8 +427,6 @@ 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,6 +143,31 @@
   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.label.startsWith("CALL")) {
+      if (d.callNode) {
         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.label.startsWith("CALL\n") ? [d] : []).join("g")
+  addTags(nodes.selectAll("g.type").data(d => d.callNode ? [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");
+  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");
   for (const [k,v] of rows) { root.append("div").text(k); root.append("div").node().append(v); }
   return root;
 }
@@ -253,10 +253,19 @@ 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);
-  const html = d3.create("div").classed("info", true);
+  if (metadata.querySelector(".info") == null) d3.select(metadata).html("").append("div").classed("info", true);
+  const html = d3.select(".info").html("");
   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());
@@ -290,7 +299,24 @@ function setFocus(key) {
       if (shape != null) p.style("cursor", "pointer").on("click", () => setFocus(shape));
     }
   }
-  return metadata.replaceChildren(html.node());
+  // 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" });
+  }
 }
 
 const EventTypes = { EXEC:0, BUF:1 };
@@ -299,7 +325,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}; focusedDevice = null; focusedShape = null; }
+  if (data?.path !== path) { data = {tracks:new Map(), axes:{}, path, first:null, pcToShape:new Map()}; focusedDevice = null; focusedShape = null; }
   setFocus(focusedShape);
   // layout once!
   if (data.tracks.size !== 0) return updateProgress(Status.COMPLETE);
@@ -312,9 +338,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();
+  const dur = u32(), tracePeak = u64(), indexLen = u32(), layoutsLen = u32(); data.dur = dur;
   const textDecoder = new TextDecoder("utf-8");
-  const { strings, dtypeSize, markers }  = JSON.parse(textDecoder.decode(new Uint8Array(buf, offset, indexLen))); offset += indexLen;
+  const { strings, dtypeSize, markers, ...extData } = 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");
@@ -333,7 +359,8 @@ 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 colors = colorScheme[k.split(":")[0]] ?? colorScheme.DEFAULT;
+    const [dname, dnum] = k.split(":", 2);
+    const colors = colorScheme[dname] ?? colorScheme.DEFAULT;
     const offsetY = baseY-canvasTop+padding/2;
     const shapes = [], visible = [];
     const eventType = u8(), eventsLen = u32();
@@ -375,8 +402,10 @@ 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 stepIdx = ctxs[ref.ctx+1].steps.findIndex((s, i) => i >= start && s.name == e.name);
-          if (stepIdx !== -1) { ref.step = stepIdx; shapeRef = ref; }
+          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; }
+          }
         } else {
           const steps = ctxs[state.currentCtx].steps;
           for (let i=state.currentStep+1; i<steps.length; i++) {
@@ -388,8 +417,9 @@ 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("");
-        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 };
+        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 };
         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 });
@@ -398,65 +428,74 @@ 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;
-          const shape = {x:[x], y:[y], dtype, sz, nbytes, key};
-          buf_shapes.set(key, shape); temp.set(key, shape);
+          allocs.set(key, {nbytes, shapeKey:`${k}-${shapeIdx++}`});
+          memEvents.push({alloc, key, dtype, sz, nbytes});
           timestamps.push(ts);
           x += 1; y += nbytes; valueMap.set(ts, y);
         } else {
-          const free = buf_shapes.get(key);
-          free.users = Array.from({ length: u32() }, () => ({shape:shapeMap.get(u32()), repr:strings[u32()], num:u32(), mode:u8()}));
+          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});
           timestamps.push(ts); valueMap.set(ts, y);
-          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);
-          }
+          x += 1; y -= 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 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));
+      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 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);
-      }
+        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;
+      };
       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,
-                           views:[[sum], shapes], valueMap, rowBorderColor });
+                           get views() { return [[sum], buildBufShapes()]; }, 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;
@@ -473,6 +512,8 @@ 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;
@@ -498,10 +539,14 @@ 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
@@ -535,7 +580,7 @@ async function renderProfiler(path, unit, opts) {
       }
       // draw row line
       if (rowBorderColor != null) {
-        const y = offsetY+rect(document.getElementById(k)).height-padding/2 - 0.5;
+        const y = offsetY+trackHeight-padding/2 - 0.5;
         drawLine(ctx, [0, canvasWidth], [y, y], { color:rowBorderColor });
       }
     }
@@ -599,6 +644,7 @@ 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;
@@ -813,6 +859,8 @@ 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;
@@ -844,7 +892,6 @@ 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)];
@@ -1034,13 +1081,15 @@ 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
-  if (event.key == "ArrowLeft") {
-    event.preventDefault()
-    return setState({ currentRewrite:Math.max(0, currentRewrite-1) });
-  }
-  if (event.key == "ArrowRight") {
+  // left and right go through rewrites in a single UOp, in profiler go forward/backward in time
+  if (event.key == "ArrowLeft" || 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) });
   }