Merge branch 'master' into sym_work

* work

* 2

* better color

.github/workflows/autogen.yml

@@ -32,6 +32,7 @@ jobs:
     - name: Setup Environment
       uses: ./.github/actions/setup-tinygrad
       with:
+        key: 'autogen'
         opencl: 'true'
         amd: 'true'
         cuda: 'true'
@@ -81,6 +82,7 @@ jobs:
     - name: Setup Environment
       uses: ./.github/actions/setup-tinygrad
       with:
+        key: 'autogen-mac'
         llvm: 'true'
     - name: Regenerate autogen files
       run: |
@@ -110,6 +112,8 @@ jobs:
       uses: actions/checkout@v4
     - name: Setup Environment
       uses: ./.github/actions/setup-tinygrad
+      with:
+        key: 'autogen-comgr'
     - name: Install autogen support packages
       run: |
         wget https://repo.radeon.com/rocm/rocm.gpg.key -O - | gpg --dearmor | sudo tee /etc/apt/keyrings/rocm.gpg > /dev/null

.github/workflows/benchmark.yml

@@ -21,6 +21,9 @@ jobs:
   # the 3 minute timeout should not be raised
   testmacpytest:
     name: Mac pytest
+    env:
+      CI: ""
+      CAPTURE_PROCESS_REPLAY: "0"
     runs-on: [self-hosted, macOS]
     timeout-minutes: 3
     defaults:
@@ -41,22 +44,14 @@ jobs:
       run: |
         echo "CACHEDB=/tmp/pytest-db-ci.db" >> $GITHUB_ENV
         rm -f /tmp/pytest-db-ci*
-    # TODO: remove this step once all old caches are migrated
-    - name: Migrate old huggingface cache (symlinks break onnxruntime 1.24+)
-      run: |
-        cd ~/Library/Caches/tinygrad/downloads/models 2>/dev/null || exit 0
-        for old_dir in models--*; do
-          [ -d "$old_dir" ] || continue
-          repo_id=$(echo "$old_dir" | sed 's/models--//; s/--/\//g')
-          snapshot=$(ls -1 "$old_dir/snapshots" 2>/dev/null | head -1)
-          [ -n "$snapshot" ] || continue
-          mkdir -p "$repo_id"
-          cp -RLn "$old_dir/snapshots/$snapshot/"* "$repo_id/" 2>/dev/null || true
-        done
     - name: Run pytest -nauto
       run: |
         source /tmp/tinygrad_pytest_ci/bin/activate
         pytest -nauto --durations=20
+    - name: openpilot compile3 0.10.1 driving_vision
+      run: FLOAT16=1 CL=1 IMAGE=2 python3.11 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
+    - name: IMAGE=1 openpilot compile3 0.10.1 driving_vision
+      run: FLOAT16=1 CL=1 IMAGE=1 python3.11 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
 
   testmacbenchmark:
     name: Mac Benchmark
@@ -343,7 +338,7 @@ jobs:
     - name: Run 10 CIFAR training steps
       run: BENCHMARK_LOG=cifar_10steps ASSERT_MIN_STEP_TIME=120 NV=1 STEPS=10 python3 examples/hlb_cifar10.py
     - name: Run 10 CIFAR training steps w HALF
-      run: BENCHMARK_LOG=cifar_10steps_half ASSERT_MIN_STEP_TIME=110 NV=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py
+      run: BENCHMARK_LOG=cifar_10steps_half ASSERT_MIN_STEP_TIME=120 NV=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py
     - name: Run 10 CIFAR training steps w BF16
       run: BENCHMARK_LOG=cifar_10steps_bf16 ASSERT_MIN_STEP_TIME=120 NV=1 STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3 examples/hlb_cifar10.py
     # - name: Run 10 CIFAR training steps w winograd
@@ -515,7 +510,7 @@ jobs:
     - name: Run 10 CIFAR training steps
       run: BENCHMARK_LOG=cifar_10steps ASSERT_MIN_STEP_TIME=200 AMD=1 STEPS=10 python3 examples/hlb_cifar10.py
     - name: Run 10 CIFAR training steps w HALF
-      run: BENCHMARK_LOG=cifar_10steps_half ASSERT_MIN_STEP_TIME=200 AMD=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py
+      run: BENCHMARK_LOG=cifar_10steps_half ASSERT_MIN_STEP_TIME=230 AMD=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py
     # - name: Run 10 CIFAR training steps w BF16
     #   run: BENCHMARK_LOG=cifar_10steps_bf16 ASSERT_MIN_STEP_TIME=288 AMD=1 STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3 examples/hlb_cifar10.py
     # TODO: too slow

.github/workflows/test.yml

@@ -649,10 +649,8 @@ jobs:
         run: AMD_LLVM=0 python -m pytest -n=auto test/amd/ --durations 20
       - name: Run AMD renderer tests (AMD_LLVM=1)
         run: AMD_LLVM=1 python -m pytest -n=auto test/amd/ --durations 20
-      - name: Run TestOps.test_add with SQTT
-        run: |
-          VIZ=-2 DEBUG=5 python3 test/backend/test_ops.py TestOps.test_add
-          extra/sqtt/rgptool.py create "/tmp/profile.pkl.$USER" -o /tmp/gpu0.rgp
+      - name: Run SQTT profiling tests
+        run: PROFILE=1 SQTT=1 python3 -m pytest -n=auto test/amd/test_sqtt_profiler.py
       - name: Run AMD emulated tests on NULL backend
         env:
           AMD: 0
@@ -664,6 +662,30 @@ jobs:
       - name: Run LLVM test
         run: AMD_LLVM=1 python test/device/test_amd_llvm.py
 
+  testmockam:
+    name: Linux (am)
+    runs-on: ubuntu-24.04
+    timeout-minutes: 15
+    env:
+      AMD: 1
+      MOCKGPU: 1
+      AMD_IFACE: PCI
+    steps:
+      - name: Checkout Code
+        uses: actions/checkout@v4
+      - name: Setup Environment
+        uses: ./.github/actions/setup-tinygrad
+        with:
+          key: mockam
+          deps: testing_unit
+          amd: 'true'
+      - name: Run test_tiny on MOCKAM
+        run: python test/test_tiny.py
+      - name: Run test_tiny on MOCKAM USB
+        run: AMD_IFACE=USB python test/test_tiny.py
+      - name: Run test_hcq on MOCKAM
+        run: python -m pytest test/device/test_hcq.py
+
   testamd:
     strategy:
       fail-fast: false
@@ -802,6 +824,8 @@ jobs:
       run: METAL=1 DEBUG=3 python test/backend/test_ops.py TestOps.test_big_gemm
     - name: Test Beam Search
       run: METAL=1 IGNORE_BEAM_CACHE=1 python3 -m pytest extra/optimization/test_beam_search.py
+    - name: Test Device Specific
+      run: METAL=1 python3 -m pytest test/device/test_metal.py
     #- name: Fuzz Test linearizer
     #  run: METAL=1 DEPTH=4 FUZZ_N=50 FUZZ_MAX_SIZE=1000000 python test/external/fuzz_linearizer.py
     - name: Run TRANSCENDENTAL math

.gitignore

@@ -66,3 +66,5 @@ target
 .mypy_cache
 mutants
 .mutmut-cache
+dagre/
+graphlib/

docs/developer/layout.md

@@ -10,7 +10,7 @@ Directories are listed in order of how they are processed.
 
 Group UOps into kernels.
 
-::: tinygrad.schedule.rangeify.get_rangeify_map
+::: tinygrad.schedule.rangeify.get_kernel_graph
     options:
         members: false
         show_labels: false

examples/mlperf/dataloader.py

@@ -254,8 +254,8 @@ def load_unet3d_data(preprocessed_dataset_dir, seed, queue_in, queue_out, X:Tens
       x = random_brightness_augmentation(x)
       x = gaussian_noise(x)
 
-    X[idx].contiguous().realize().uop.base.realized.as_memoryview(force_zero_copy=True)[:] = x.tobytes()
-    Y[idx].contiguous().realize().uop.base.realized.as_memoryview(force_zero_copy=True)[:] = y.tobytes()
+    X[idx].flatten().assign(x.tobytes())
+    Y[idx].flatten().assign(y.tobytes())
 
     queue_out.put(idx)
   queue_out.put(None)
@@ -369,12 +369,12 @@ def load_retinanet_data(base_dir:Path, val:bool, queue_in:Queue, queue_out:Queue
       clipped_match_idxs = np.clip(match_idxs, 0, None)
       clipped_boxes, clipped_labels = tgt["boxes"][clipped_match_idxs], tgt["labels"][clipped_match_idxs]
 
-      boxes[idx].contiguous().realize().uop.base.realized.as_memoryview(force_zero_copy=True)[:] = clipped_boxes.tobytes()
-      labels[idx].contiguous().realize().uop.base.realized.as_memoryview(force_zero_copy=True)[:] = clipped_labels.tobytes()
-      matches[idx].contiguous().realize().uop.base.realized.as_memoryview(force_zero_copy=True)[:] = match_idxs.tobytes()
-      anchors[idx].contiguous().realize().uop.base.realized.as_memoryview(force_zero_copy=True)[:] = anchor.tobytes()
+      boxes[idx].flatten().assign(clipped_boxes.tobytes())
+      labels[idx].flatten().assign(clipped_labels.tobytes())
+      matches[idx].flatten().assign(match_idxs.tobytes())
+      anchors[idx].flatten().assign(anchor.tobytes())
 
-    imgs[idx].contiguous().realize().uop.base.realized.as_memoryview(force_zero_copy=True)[:] = img.tobytes()
+    imgs[idx].flatten().assign(img.tobytes())
 
     queue_out.put(idx)
   queue_out.put(None)

examples/mlperf/model_train.py

@@ -1335,6 +1335,9 @@ def train_llama3():
   model_params = MODEL_PARAMS[getenv("LLAMA3_SIZE", "8B")]["args"]
   # vocab_size from the mixtral tokenizer
   if not SMALL: model_params |= {"vocab_size": 32000}
+  real_vocab_size = model_params['vocab_size']
+  if (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}")
 
@@ -1352,6 +1355,8 @@ def train_llama3():
     for v in get_parameters(model):
       v.shard_(device, axis=None)
 
+    vocab_mask.shard_(device, axis=None)
+
   if (MP := getenv("MP", 1)) > 1:
     device = tuple(f"{Device.DEFAULT}:{i}" for i in range(MP))
     for k,v in get_state_dict(model).items():
@@ -1359,6 +1364,7 @@ def train_llama3():
       elif '.attention.wq' in k: v.shard_(device, axis=0)
       elif '.attention.wk' in k: v.shard_(device, axis=0)
       elif '.attention.wv' in k: v.shard_(device, axis=0)
+      elif '.attention.wqkv' in k: v.shard_(device, axis=0)
       elif '.attention.wo' in k: v.shard_(device, axis=1)
       elif '.feed_forward.w1.' in k: v.shard_(device, axis=0)
       elif '.feed_forward.w2.' in k: v.shard_(device, axis=1)
@@ -1371,8 +1377,11 @@ def train_llama3():
       # prevents memory spike on device 0
       v.realize()
 
-  optim = GradAccClipAdamW(get_parameters(model), lr=0.0,
-                b1=opt_adamw_beta_1, b2=opt_adamw_beta_2, eps=opt_adamw_epsilon, weight_decay=opt_adamw_weight_decay, grad_acc=grad_acc)
+    vocab_mask.shard_(device, axis=2).realize()
+
+  optim_device = "CPU" if getenv("OFFLOAD_OPTIM") else None
+  optim = GradAccClipAdamW(get_parameters(model), lr=0.0, b1=opt_adamw_beta_1, b2=opt_adamw_beta_2,
+                           eps=opt_adamw_epsilon, weight_decay=opt_adamw_weight_decay, grad_acc=grad_acc, device=optim_device)
 
   # init grads
   for p in optim.params:
@@ -1400,7 +1409,7 @@ def train_llama3():
       device = tuple(f"{Device.DEFAULT}:{i}" for i in range(MP))
       tokens = tokens.shard(device)
     logits:Tensor = model(tokens[:, :-1], start_pos=0, temperature=math.nan)
-    loss = logits.sparse_categorical_crossentropy(tokens[:, 1:])
+    loss = vocab_mask.where(-float("inf"), logits).sparse_categorical_crossentropy(tokens[:, 1:])
     loss.backward()
     assert all(p.grad is g for p,g in zip(optim.params, grads))
     Tensor.realize(loss, *grads)
@@ -1430,7 +1439,7 @@ def train_llama3():
       device = tuple(f"{Device.DEFAULT}:{i}" for i in range(MP))
       tokens = tokens.shard(device)
     logits:Tensor = model(tokens[:, :-1], start_pos=0, temperature=math.nan)
-    loss = logits.sparse_categorical_crossentropy(tokens[:, 1:])
+    loss = vocab_mask.where(-float("inf"), logits).sparse_categorical_crossentropy(tokens[:, 1:])
     return loss.flatten().float().to("CPU")
 
   # ** data iters **
@@ -1468,29 +1477,28 @@ def train_llama3():
       st = time.perf_counter()
 
       stopped = False
+      losses, data_time, dev_time = [], 0, 0
       for _ in range(grad_acc):
         ist = time.perf_counter()
         try: tokens = next(train_iter)
         except StopIteration:
           stopped = True
           break
-        dt = time.perf_counter()
-        loss = minibatch(tokens)
+        mst = time.perf_counter()
+        data_time += mst - ist
+        losses.append(minibatch(tokens).item())
+        dev_time += time.perf_counter() - mst
       if stopped: break
 
       gt = time.perf_counter()
-      lr = optim_step()
-      ot = time.perf_counter()
-
-      loss = loss.float().item()
-      lr = lr.item()
-
+      lr = optim_step().item()
       et = time.perf_counter()
+
+      loss = sum(losses) / len(losses)
+      optim_time = et - gt
+      dev_time += optim_time
       step_time = et - st
       gbs_time = gt - st
-      optim_time = ot - gt
-      data_time = dt - ist
-      dev_time = step_time - data_time * grad_acc
       if BENCHMARK: step_times.append(step_time)
 
       i += 1

examples/mlperf/optim.py

@@ -1,24 +1,47 @@
 from tinygrad.tensor import Tensor
 from tinygrad.dtype import dtypes
-from tinygrad.nn.optim import LAMB
+from tinygrad.nn.optim import Optimizer
 from tinygrad.helpers import FUSE_OPTIM
 
-class GradAccClipAdamW(LAMB):
-  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, fused=FUSE_OPTIM):
-    super().__init__(params, lr, b1, b2, eps, weight_decay, adam=True, fused=FUSE_OPTIM)
+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.m = self._new_optim_param()
+    self.v = self._new_optim_param()
     self.grad_acc, self.clip_norm = grad_acc, clip_norm
 
   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 self.fused:
       grads[0] = 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)
     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()
       for i in range(len(grads)):
-        grads[i] = grads[i] / self.grad_acc
         grads[i] = (grads[i] * (self.clip_norm / (total_norm + 1e-6)).clamp(max_=1.0)).cast(grads[i].dtype)
-    return super()._step(params, grads)
+
+    ret = []
+    self.b1_t *= self.b1
+    self.b2_t *= self.b2
+    for i, (t, g) in enumerate(zip(params, grads)):
+      self.m[i].assign((self.b1 * self.m[i] + (1.0 - self.b1) * g).cast(self.m[i].dtype))
+      self.v[i].assign((self.b2 * self.v[i] + (1.0 - self.b2) * (g * g)).cast(self.v[i].dtype))
+      m_hat = self.m[i] / (1.0 - self.b1_t)
+      v_hat = self.v[i] / (1.0 - self.b2_t)
+      up = m_hat / (v_hat.sqrt() + self.eps)
+      ret.append((self.lr * up).cast(t.dtype))
+    return ret, [self.b1_t, self.b2_t] + self.m + self.v
+
+  def _apply_update(self, t:Tensor, up:Tensor) -> Tensor:
+    up = up.shard_like(t) + self.lr.to(t.device) * self.wd * t.detach()
+    return t.detach() - up.cast(t.dtype)

extra/gemm/asm/cdna/gemm.py

@@ -11,12 +11,12 @@ from extra.gemm.asm.cdna.asm import build_kernel, TILE_M, TILE_N, TILE_K, NUM_WG
 WORKGROUP_SIZE = 256
 
 @functools.cache
-def custom_asm_gemm(C:UOp, A:UOp, B:UOp, dname:str, arch:str, wg:int) -> UOp:
+def custom_asm_gemm(C:UOp, A:UOp, B:UOp, dname:str) -> UOp:
   batch, M, K = A.shape
   K2, N = B.shape[(1 if B.ndim == 3 else 0):]
   assert K == K2
   lidx = UOp.special(WORKGROUP_SIZE, "lidx0")
-  gidx = UOp.special(wg, "gidx0")
+  gidx = UOp.special(NUM_WG, "gidx0")
   insts = build_kernel(batch, M, N, K, A.dtype.base)
   lds = UOp(Ops.DEFINE_LOCAL, dtypes.uint8.ptr(size=133_120, addrspace=AddrSpace.LOCAL), (), 'lds')
   sink = UOp.sink(C.base, A.base, B.base, lds, lidx, gidx,
@@ -94,7 +94,7 @@ def asm_gemm(a:Tensor, b:Tensor) -> Tensor:
   renderer = Device[a.device[0] if is_multi else a.device].renderer
   dname, arch = renderer.device, getattr(renderer, "arch", "")
   if arch.startswith("gfx950") and getenv("USE_ASM", 1):
-    out = Tensor.custom_kernel(out, a, b, fxn=functools.partial(custom_asm_gemm, dname=dname, wg=NUM_WG, arch=arch), grad_fxn=custom_gemm_bw)[0]
+    out = Tensor.custom_kernel(out, a, b, fxn=functools.partial(custom_asm_gemm, dname=dname), grad_fxn=custom_gemm_bw)[0]
   else:
     out = Tensor.custom_kernel(out, a, b, fxn=custom_uop_gemm, grad_fxn=custom_gemm_bw)[0]
   if k_sharded: out = out.sum(0)

extra/models/llama.py

@@ -56,7 +56,10 @@ class Attention:
   def __call__(self, x:Tensor, start_pos:Union[Variable,int], freqs_cis:Tensor, mask:Optional[Tensor]=None) -> Tensor:
     if getenv("WQKV"):
       xqkv = self.wqkv(x)
-      xq, xk, xv = xqkv.split([self.n_heads * self.head_dim, self.n_kv_heads * self.head_dim, self.n_kv_heads * self.head_dim], dim=2)
+      xqkv = xqkv.reshape(xqkv.shape[0], xqkv.shape[1], self.n_kv_heads, self.n_rep + 2, self.head_dim)
+      xq = xqkv[:, :, :, :self.n_rep].reshape(xqkv.shape[0], xqkv.shape[1], -1)
+      xk = xqkv[:, :, :, self.n_rep:self.n_rep+1].reshape(xqkv.shape[0], xqkv.shape[1], -1)
+      xv = xqkv[:, :, :, self.n_rep+1:self.n_rep+2].reshape(xqkv.shape[0], xqkv.shape[1], -1)
     else:
       xq, xk, xv = self.wq(x), self.wk(x.contiguous_backward()), self.wv(x)
 

extra/sqtt/examples/generate_examples.py

@@ -1,4 +1,4 @@
-import os, subprocess
+import os, subprocess, sys
 from pathlib import Path
 from tinygrad.helpers import temp
 
@@ -6,9 +6,9 @@ EXAMPLES_DIR = Path(__file__).parent
 PROFILE_PATH = Path(temp("profile.pkl", append_user=True))
 
 EXAMPLES = [
-  "test.backend.test_custom_kernel.TestCustomKernel.test_empty",
-  "test.test_tiny.TestTiny.test_plus",
-  "test.test_tiny.TestTiny.test_gemm",
+  "test/backend/test_custom_kernel.py TestCustomKernel.test_empty",
+  "test/test_tiny.py TestTiny.test_plus",
+  "test/test_tiny.py TestTiny.test_gemm",
 ]
 
 if __name__ == "__main__":
@@ -17,7 +17,8 @@ if __name__ == "__main__":
   (EXAMPLES_DIR/arch).mkdir(exist_ok=True)
   for test in EXAMPLES:
     for i in range(2):
-      subprocess.run(["python", "-m", "unittest", test], cwd=EXAMPLES_DIR.parent.parent.parent,
-                     env={**os.environ, "AMD":"1", "SQTT_LIMIT_SE":"-1", "VIZ":"-2"}, check=True)
+      # AM_RESET=1 gets a clear trace, does not work on mi300 machines
+      subprocess.run([sys.executable, *test.split()], cwd=EXAMPLES_DIR.parent.parent.parent,
+                     env={**os.environ, "AMD":"1", "AM_RESET":"1" if not arch.startswith("gfx9") else "0", "VIZ":"-2", "PYTHONPATH":"."})
       PROFILE_PATH.rename(dest:=EXAMPLES_DIR/arch/f"profile_{test.split('.')[-1].replace('test_', '')}_run_{i}.pkl")
       print(f"saved SQTT trace to {dest}")

extra/thunder/amd/fa.py

@@ -11,7 +11,8 @@ from tinygrad.uop.ops import UOp, Ops, KernelInfo
 def _sharded_empty(shape:Tensor, ref:Tensor, axis:int|None, dtype:DTypeLike|None=None) -> Tensor:
   dtype = dtype or ref.dtype
   if not isinstance(ref.device, tuple): return Tensor.empty(*shape, dtype=dtype, device=ref.device)
-  shape = tuple(s // len(ref.device) if i == ref.uop.axis else s for i, s in enumerate(shape))
+  shard_axis = ref.uop.axis if axis is None else axis
+  shape = tuple(s // len(ref.device) if i == shard_axis else s for i, s in enumerate(shape))
   axis = ref.uop.axis if axis is None else axis
   return Tensor(Tensor.empty(*shape, dtype=dtype, device=ref.device).uop.multi(axis), dtype=dtype, device=ref.device)
 
@@ -29,34 +30,40 @@ def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False
   assert D == 128, "only D=128 supported"
 
   num_devices = len(xq.device) if isinstance(xq.device, tuple) else 1
-  B_local = B // num_devices
-  if DEBUG >= 2: print(f"Flash Attention {B=} {B_local=} {N=} {H=} {H_KV=} {D=}")
+  is_dp = xq.uop.axis == 0
+  is_mp = xq.uop.axis == 2
+  B_local = B // num_devices if is_dp else B
+  H_local = H // num_devices if is_mp else H
+  H_KV_local = H_KV // num_devices if is_mp else H_KV
+  shard_axis = 0 if is_dp else 2 if is_mp else None
+  shard_axis_t = 0 if is_dp else 1 if is_mp else None
+  if DEBUG >= 2: print(f"Flash Attention {B=} {B_local=} {N=} {H=} {H_local=} {H_KV=} {H_KV_local=} {D=} on {num_devices} devices, {'DP' if is_dp else 'MP' if is_mp else 'no sharding'}")
 
   single_device = xq.device[0] if isinstance(xq.device, tuple) else xq.device
   arch = Device[single_device].renderer.arch
 
-  attn = _sharded_empty_like(xq, axis=0)
-  l_vec = _sharded_empty((B, H, 1, N), xq, axis=0, dtype=dtypes.float32)
+  attn = _sharded_empty_like(xq, axis=shard_axis)
+  l_vec = _sharded_empty((B, H, 1, N), xq, dtype=dtypes.float32, axis=shard_axis_t)
 
   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=0)
-    dq = _sharded_empty_like(xq, axis=0)
-    dk = _sharded_empty_like(xk, axis=0)
-    dv = _sharded_empty_like(xv, axis=0)
+    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)
 
     # delta_vec = (do * attn).sum(-1, dtype=dtypes.float32).transpose(1, 2).unsqueeze(-2).detach()
-    delta_vec = _sharded_empty((B, H, 1, N), xq, axis=0, dtype=dtypes.float32)
-    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, H_KV=H_KV, D=D))[:2]
+    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]
 
-    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, H_KV=H_KV, D=D))[:3]
+    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]
 
     # 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, H_KV=H_KV, D=D))[0]
+    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]
 
     return None, None, dq.uop, dk.uop, dv.uop
 
-  attn, l_vec = Tensor.custom_kernel(attn, l_vec, xq, xk, xv, fxn=functools.partial(custom_fa_forward, device=single_device, arch=arch, B=B_local, N=N, H=H, H_KV=H_KV, D=D), grad_fxn=grad)[:2]
+  attn, l_vec = Tensor.custom_kernel(attn, l_vec, xq, xk, xv, fxn=functools.partial(custom_fa_forward, device=single_device, arch=arch, B=B_local, N=N, H=H_local, H_KV=H_KV_local, D=D), grad_fxn=grad)[:2]
 
   return attn.transpose(1, 2)
 

extra/torch_backend/example.py

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

extra/torch_backend/test_kernel_fusion.py

@@ -26,7 +26,7 @@ class TestKernelFusionRegression(unittest.TestCase):
     def fn():
       x = torch.randn(128, 128, device=device)
       return (x + 1.0) * 2.0 - 0.5
-    self._check_kernel_count(fn, 6)
+    self._check_kernel_count(fn, 5)
 
   def test_relu_fusion(self):
     def fn():
@@ -50,14 +50,14 @@ class TestKernelFusionRegression(unittest.TestCase):
     def fn():
       x = torch.randn(64, 64, device=device)
       return (x * 2.0).sum()
-    self._check_kernel_count(fn, 7)
+    self._check_kernel_count(fn, 5)
 
   def test_matmul_elementwise_fusion(self):
     def fn():
       x = torch.randn(32, 32, device=device)
       w = torch.randn(32, 32, device=device)
       return torch.nn.functional.relu(x @ w + 1.0)
-    self._check_kernel_count(fn, 6)
+    self._check_kernel_count(fn, 7)
 
   def test_pooling_fusion(self):
     def fn():
@@ -71,7 +71,7 @@ class TestKernelFusionRegression(unittest.TestCase):
       identity = torch.randn(1, 8, 16, 16, device=device)
       out = x + identity
       return torch.nn.functional.relu(out)
-    self._check_kernel_count(fn, 6)
+    self._check_kernel_count(fn, 7)
 
   def test_inplace_add_relu_fusion(self):
     def fn():
@@ -79,7 +79,7 @@ class TestKernelFusionRegression(unittest.TestCase):
       y = torch.randn(1, 16, 32, 32, device=device)
       x += y
       return torch.nn.functional.relu(x)
-    self._check_kernel_count(fn, 6)
+    self._check_kernel_count(fn, 7)
 
   def test_conv_bn_add_relu_fusion(self):
     def fn():
@@ -92,7 +92,7 @@ class TestKernelFusionRegression(unittest.TestCase):
         out = bn(conv(x))
         out += identity
         return torch.nn.functional.relu(out)
-    self._check_kernel_count(fn, 16)
+    self._check_kernel_count(fn, 17)
 
   def test_multiple_inplace_ops_fusion(self):
     def fn():
@@ -138,7 +138,7 @@ class TestKernelFusionRegression(unittest.TestCase):
       loss.backward()
       optimizer.step()
       return loss
-    self._check_kernel_count(fn, 33)
+    self._check_kernel_count(fn, 28)
 
 if __name__ == "__main__":
   unittest.main()

test/amd/test_sqtt_examples.py

@@ -208,12 +208,12 @@ class SQTTExamplesTestBase(unittest.TestCase):
 class TestSQTTExamplesRDNA3(SQTTExamplesTestBase):
   target = "gfx1100"
   expected = {
-    "profile_empty_run_0": [1844, 1885, 1905, 1956, 1983, 1889],
-    "profile_empty_run_1": [1780, 1885, 1905, 1956, 1983, 1889],
-    "profile_gemm_run_0": [2656, 2025, 2045, 2096, 2123, 2029, 3183, 2019, 2039, 2090, 2117, 2023, 19119, 2013, 2033, 2084, 2111, 2017],
-    "profile_gemm_run_1": [2662, 2025, 2045, 2096, 2123, 2029, 3179, 2019, 2039, 2090, 2117, 2023, 19113, 2071, 2091, 2142, 2169, 2075],
-    "profile_plus_run_0": [1886, 2013, 2033, 2084, 2111, 2017],
-    "profile_plus_run_1": [1988, 2071, 2091, 2142, 2169, 2075],
+    "profile_empty_run_0": [1744, 1801, 1854, 1890, 1917, 1822],
+    "profile_empty_run_1": [1744, 1801, 1854, 1886, 1921, 1906],
+    "profile_gemm_run_0": [1800, 1867, 1899, 1898, 1914, 1895, 1694, 1779, 1819, 1872, 1877, 1858, 1750, 1834, 1866, 1834, 1911, 1796],
+    "profile_gemm_run_1": [1806, 1874, 1837, 1885, 1907, 1906, 1694, 1778, 1810, 1873, 1885, 1867, 1750, 1834, 1866, 1856, 1903, 1897],
+    "profile_plus_run_0": [1744, 1878, 1854, 1890, 1878, 1910],
+    "profile_plus_run_1": [1744, 1878, 1854, 1886, 1921, 1909],
   }
 
 class TestSQTTExamplesRDNA4(SQTTExamplesTestBase): target = "gfx1200"

test/amd/test_sqtt_profiler.py

@@ -0,0 +1,94 @@
+import unittest, contextlib
+from tinygrad import Device, Tensor, Context, TinyJit
+from tinygrad.device import Compiled, ProfileProgramEvent, ProfileDeviceEvent
+from tinygrad.viz.serve import load_amd_counters
+
+@contextlib.contextmanager
+def save_sqtt():
+  yield (ret:=[])
+  Device[Device.DEFAULT].synchronize()
+  Device[Device.DEFAULT]._at_profile_finalize()
+  load_amd_counters(ret, Compiled.profile_events)
+  ret[:] = [r for r in ret if r["name"].startswith("Exec")]
+
+@unittest.skipUnless(Device.DEFAULT == "AMD", "only runs on AMD")
+class TestSQTTProfiler(unittest.TestCase):
+  # TODO: can we enable SQTT profiling in context?
+  @classmethod
+  def setUpClass(cls):
+    if not Device[Device.DEFAULT].sqtt_enabled: raise unittest.SkipTest("device must be in SQTT profiling mode")
+
+  def setUp(self):
+    Device[Device.DEFAULT].synchronize()
+    Compiled.profile_events[:] = [e for e in Compiled.profile_events if isinstance(e, (ProfileProgramEvent, ProfileDeviceEvent))]
+
+  def test_simple(self):
+    t = Tensor.empty(1) + 1
+    with save_sqtt() as sqtt:
+      ei = t.schedule()[0].lower()
+      ei.run()
+    self.assertEqual(len(sqtt), 1)
+    self.assertEqual(sqtt[0]["name"], f"Exec {ei.prg.p.function_name}")
+
+  def test_multiple_runs(self):
+    t = Tensor.empty(1) + 1
+    with save_sqtt() as sqtt:
+      ei = t.schedule()[0].lower()
+      for _ in range(N:=3):
+        ei.run()
+    self.assertEqual(len(sqtt), N)
+    for i in range(1, N):
+      self.assertEqual(sqtt[i]["name"], f"Exec {ei.prg.p.function_name} n{i+1}")
+
+  def test_multiple_kernels(self):
+    t = ((Tensor.empty(1) + 1).contiguous() + 2)
+    sched = t.schedule()
+    with save_sqtt() as sqtt:
+      for si in sched: si.lower().run()
+    self.assertEqual(len(sqtt), len(sched))
+    for i,k in enumerate(sched):
+      self.assertEqual(sqtt[i]["name"], f"Exec {k.lower().prg.p.function_name}")
+
+  def test_multiple_kernels_lower(self):
+    t = ((Tensor.empty(1) + 1).contiguous() + 2)
+    sched = t.schedule()
+    with save_sqtt() as sqtt:
+      prgs = [si.lower() for si in sched]
+      for p in prgs: p.run()
+    self.assertEqual(len(sqtt), len(sched))
+    for i,ei in enumerate(prgs):
+      self.assertEqual(sqtt[i]["name"], f"Exec {ei.prg.p.function_name}")
+
+  def test_jit(self):
+    @TinyJit
+    def f(a): return a + 1
+    t = Tensor.empty(1)
+    with save_sqtt() as sqtt:
+      for _ in range(N:=5):
+        f(t).realize()
+    self.assertEqual(len(sqtt), N)
+    kernel_name = sqtt[0]["name"]
+    for i,s in enumerate(sqtt[1:], start=1): self.assertEqual(s["name"], f"{kernel_name} n{i+1}")
+
+  # TODO: can we trace SQTT for graphed kernels?
+  def test_jit_graph(self, kernel_count=3*2):
+    @TinyJit
+    def f(a): return ((a + 1).contiguous() + 2).contiguous().sum()
+    t = Tensor.empty(32)
+    with save_sqtt() as sqtt:
+      for _ in range(5):
+        f(t).realize()
+    names = [s["name"] for s in sqtt]
+    k0, k1, k2 = names[:3]
+    for i in range(3, len(sqtt), 3):
+      n = (i // 3)+1
+      self.assertEqual(names[i], f"{k0} n{n}")
+      self.assertEqual(names[i+1], f"{k1} n{n}")
+      self.assertEqual(names[i+2], f"{k2} n{n}")
+    self.assertEqual(len(sqtt), kernel_count)
+
+  @Context(JIT=2)
+  def test_jit_multiple_kernels(self): self.test_jit_graph(kernel_count=3*5)
+
+if __name__ == "__main__":
+  unittest.main()

test/backend/test_asm_gemm.py

@@ -67,6 +67,7 @@ class TestGemmLarge(unittest.TestCase):
     if not is_cdna4():
       self.skipTest("very slow on non mi350x")
 
+  def test_tiny(self): verify_asm_gemm(1, 256, 256, 64)
   def test_simple(self): verify_asm_gemm(1, N:=getenv("N", 4096), N, N, dtype=dtypes.half)
   def test_gemm(self): verify_asm_gemm(1, 8192, 4096, 14336)
   def test_gemm_batched(self): verify_asm_gemm(2, 8192, 4096, 4096)

test/backend/test_custom_kernel.py

@@ -265,8 +265,6 @@ class TestCustomKernel(unittest.TestCase):
     Expected schedule order: [A2, B2, E, custom_addmul, final_sum]
     The custom_addmul kernel should be at index 3.
     """
-    from tinygrad.engine.schedule import create_schedule
-    from tinygrad.schedule.rangeify import get_rangeify_map
 
     A, B = Tensor.empty(4, 4), Tensor.empty(4, 4)
     A2 = (A + 1).contiguous()                      # kernel 0: depends on A
@@ -275,11 +273,7 @@ class TestCustomKernel(unittest.TestCase):
     C, D, _, _ = Tensor.custom_kernel(C, D, A2, B2, fxn=custom_elementwise_addmul_kernel)  # depends on A2 AND B2
     E = (A2 * 3).contiguous()                      # kernel 2: depends only on A2
     result = (C + D + E).sum()                     # kernel 3: custom_addmul, then kernel 4: sum
-
-    big_sink = result.uop.sink()
-    tensor_map = get_rangeify_map(big_sink)
-    sched_sink = big_sink.substitute(tensor_map)
-    schedule, _ = create_schedule(sched_sink)
+    schedule = result.schedule()
 
     # Find the custom_addmul kernel position
     custom_idx = next((i for i, item in enumerate(schedule)

test/backend/test_dtype.py

@@ -150,28 +150,16 @@ class TestFp8sConversions(unittest.TestCase):
     np.testing.assert_equal(float_to_fp8(x, dtypes.fp8e4m3), torch.tensor(x, dtype=torch.float8_e4m3fn).view(torch.uint8).item())
 
   def test_float_to_fp8e4m3_extreme_values(self):
-    np.testing.assert_equal(float_to_fp8(FP8E4M3_MAX, dtypes.fp8e4m3), 126)
-    np.testing.assert_equal(float_to_fp8(FP8E4M3_MAX*1.01, dtypes.fp8e4m3), 126)
-    np.testing.assert_equal(float_to_fp8(math.inf, dtypes.fp8e4m3), 127)
-    np.testing.assert_equal(float_to_fp8(-FP8E4M3_MAX, dtypes.fp8e4m3), 254)
-    np.testing.assert_equal(float_to_fp8(-FP8E4M3_MAX*1.01, dtypes.fp8e4m3), 254)
-    np.testing.assert_equal(float_to_fp8(-math.inf, dtypes.fp8e4m3), 255)
-    np.testing.assert_equal(float_to_fp8(math.nan, dtypes.fp8e4m3), 127)
-    np.testing.assert_equal(float_to_fp8(-math.nan, dtypes.fp8e4m3), 255)
+    for x in [FP8E4M3_MAX, FP8E4M3_MAX*1.01, -FP8E4M3_MAX, -FP8E4M3_MAX*1.01, math.inf, -math.inf, math.nan, -math.nan]:
+      np.testing.assert_equal(float_to_fp8(x, dtypes.fp8e4m3), torch.tensor(x, dtype=torch.float8_e4m3fn).view(torch.uint8).item())
 
   @given(strat.floats(width=32, allow_subnormal=True, allow_nan=False, allow_infinity=False, min_value=-FP8E5M2_MAX, max_value=FP8E5M2_MAX))
   def test_float_to_fp8e5m2(self, x):
     np.testing.assert_equal(float_to_fp8(x, dtypes.fp8e5m2), torch.tensor(x, dtype=torch.float8_e5m2).view(torch.uint8).item())
 
   def test_float_to_fp8e5m2_extreme_values(self):
-    np.testing.assert_equal(float_to_fp8(FP8E5M2_MAX, dtypes.fp8e5m2), 123)
-    np.testing.assert_equal(float_to_fp8(FP8E5M2_MAX*1.01, dtypes.fp8e5m2), 123)
-    np.testing.assert_equal(float_to_fp8(math.inf, dtypes.fp8e5m2), 124)
-    np.testing.assert_equal(float_to_fp8(-FP8E5M2_MAX, dtypes.fp8e5m2), 251)
-    np.testing.assert_equal(float_to_fp8(-FP8E5M2_MAX*1.01, dtypes.fp8e5m2), 251)
-    np.testing.assert_equal(float_to_fp8(-math.inf, dtypes.fp8e5m2), 252)
-    np.testing.assert_equal(float_to_fp8(math.nan, dtypes.fp8e5m2), 126)
-    np.testing.assert_equal(float_to_fp8(-math.nan, dtypes.fp8e5m2), 254)
+    for x in [FP8E5M2_MAX, FP8E5M2_MAX*1.01, -FP8E5M2_MAX, -FP8E5M2_MAX*1.01, math.inf, -math.inf, math.nan, -math.nan]:
+      np.testing.assert_equal(float_to_fp8(x, dtypes.fp8e5m2), torch.tensor(x, dtype=torch.float8_e5m2).view(torch.uint8).item())
 
   @given(strat.integers(min_value=0, max_value=255))
   def test_fp8e4m3_to_float(self, x):

test/backend/test_image_dtype.py

@@ -115,7 +115,7 @@ class TestImageDType(unittest.TestCase):
     tst = data.numpy()
     it = data.cast(dtypes.imagef((9,27,4))).realize()
     # the underlying UOp is identical
-    self.assertIs(it.uop.base.realized, data.uop.base.realized)
+    #self.assertIs(it.uop.base.realized, data.uop.base.realized)
     np.testing.assert_equal(tst, it.numpy())
 
   def test_image_and_back_wrong_shape(self):

test/backend/test_jit.py

@@ -332,7 +332,6 @@ class TestJit(unittest.TestCase):
     assert len(res3) == 10, "All values should be different, rand works in jit."
     assert res3 != res2, "Jit rand is diff with diff seeds"
 
-  #@unittest.expectedFailure # requires contiguous folding
   def test_jit_random_after_unrealized_random(self):
     @TinyJit
     def f(): return Tensor.rand()
@@ -476,7 +475,7 @@ class TestJit(unittest.TestCase):
     b = f(Tensor([2.0]))
     assert abs((a - b).item()) > 0.5
 
-  def test_jit_init_with_empty_different_size(self):
+  def test_jit_init_empty(self):
     @TinyJit
     def f(x:Tensor) -> Tensor: return (x + 1).realize()
 
@@ -485,10 +484,16 @@ class TestJit(unittest.TestCase):
     # scalar const input is not allowed
     with self.assertRaises(JitError):
       f(Tensor(2.0)).item()
-    # list input has different view structure than empty(1)
-    # but okay if it's realized
-    #with self.assertRaises(JitError):
-    #  f(Tensor([2.0])).item()
+    # self.assertEqual(f(Tensor([2.0])).item(), 1.0) # TODO: wrong output, should be 3.0. currently depends on empty value
+
+  def test_jit_init_empty_alt(self):
+    @TinyJit
+    def f(a:Tensor, b:Tensor) -> Tensor: return b.assign(a+1)
+    for i in range(4):
+      a = Tensor([i])
+      b = Tensor.empty_like(a)
+      c = f(a, b)
+      self.assertEqual(c.item(), i+1)
 
 @unittest.skip("Pending multioutput implementation #3607")
 class TestMultioutputJit(unittest.TestCase):

test/backend/test_multitensor.py

@@ -135,34 +135,6 @@ class TestMultiTensor(unittest.TestCase):
       si.run()
     self.assertEqual(len(set(names)), 1, "function was relinearized")
 
-  @unittest.skip("this doesn't fold because shard_ calls contiguous on all lbs")
-  def test_sharded_memory(self):
-    # Buffer may be stuck in track_cross_buffer
-    for x in (d0, d1, d2, d3, d4): Device[x].synchronize()
-    mem_base = GlobalCounters.mem_used
-
-    X = Tensor.ones(256).contiguous().realize()
-    assert GlobalCounters.mem_used-mem_base== X.dtype.itemsize * 256, GlobalCounters.mem_used-mem_base
-    X.shard_(devices_4).realize()
-    for x in (d0, d1, d2, d3, d4): Device[x].synchronize()
-    assert GlobalCounters.mem_used-mem_base == X.dtype.itemsize * 256 * 4, GlobalCounters.mem_used-mem_base
-
-    X = Tensor.ones(256).contiguous().realize()
-    assert GlobalCounters.mem_used-mem_base == X.dtype.itemsize * 256, GlobalCounters.mem_used-mem_base
-    X.shard_(devices_4, axis=0).realize()
-    for x in (d0, d1, d2, d3, d4): Device[x].synchronize()
-    assert GlobalCounters.mem_used-mem_base == X.dtype.itemsize * 256, GlobalCounters.mem_used-mem_base
-
-    X = Tensor.ones(256).realize()
-    assert GlobalCounters.mem_used-mem_base == 0
-    X.shard_(devices_4).realize()
-    assert GlobalCounters.mem_used-mem_base == 0
-
-    X = Tensor.ones(256).realize()
-    assert GlobalCounters.mem_used-mem_base == 0
-    X.shard_(devices_4, axis=0).realize()
-    assert GlobalCounters.mem_used-mem_base == 0
-
   def test_shard_same_device(self):
     X = Tensor.ones(256).contiguous().realize()
     X.shard_((d1, X.device), 0)
@@ -704,7 +676,7 @@ class TestMultiTensor(unittest.TestCase):
 
     # test no left join
     with self.assertRaises((AssertionError, ValueError)):
-      t0.reshape((26*15,7)).schedule()
+      t0.reshape((26*15,7)).contiguous().schedule()
 
   # it doesn't work like this anymore
   # NOTE: this never failed in assign_multi, it failed tensor spec because MULTI was never pushed in the graph
@@ -840,13 +812,15 @@ class TestMultiTensor(unittest.TestCase):
     t.shard_(devices, axis=0).realize()
     assert all([lb is lb.base and lb.realized.base.size == 4 * 16 for lb in t.uop.src])
 
-  @unittest.skip("this is unreliable on OSX")
   def test_clone(self):
-    t = Tensor.rand(16, 16).shard(devices_2, axis=None)
-    np.testing.assert_allclose(t.numpy(), t.clone().numpy())
-
-    t = Tensor.rand(16, 16).shard(devices_2, axis=0)
-    np.testing.assert_allclose(t.numpy(), t.clone().numpy())
+    for axis in (None, 0):
+      t = Tensor.arange(16).reshape(4, 4).shard(devices_2, axis=axis).contiguous().realize()
+      t_clone = t.clone().realize()
+      self.assertEqual(t_clone.device, t.device)
+      self.assertEqual(t_clone.uop.axis, axis)
+      self.assertEqual(t_clone.tolist(), t.tolist())
+      t_clone += 1
+      self.assertNotEqual(t_clone.tolist(), t.tolist())
 
   @unittest.skip("RANGEIFY doesn't support multi const folding")
   def test_multi_const_folding(self):
@@ -895,18 +869,18 @@ class TestShrinkMultiTensorShardedAxis(unittest.TestCase):
 
     with self.assertRaises(AssertionError):
       # sharded axis shrink on non-device boundry is not allowed
-      a = t.shrink(((0, 3), (0, 8)))
+      a = t.shrink(((0, 3), (0, 8))).contiguous()
       a.schedule()
     a = t.shrink(((0, 2), (2, 4)))
     assert a.shape == (2, 2)
     ref = Tensor.arange(64).reshape(8, 8).shrink(((0, 2), (2, 4)))
     np.testing.assert_equal(a.numpy(), ref.numpy())
 
-    a = t.shrink(((0, 2), (0, 8)))
+    a = t.shrink(((0, 2), (0, 8))).contiguous()
     a.schedule()
     assert a.shape == (2, 8)
 
-    p = a.pad(((0, 6), (0, 0)))
+    p = a.pad(((0, 6), (0, 0))).contiguous()
     p.schedule()
     assert p.shape == (8, 8)
 

test/backend/test_ops.py

@@ -8,7 +8,8 @@ from tinygrad.tensor import _to_np_dtype
 from tinygrad.device import is_dtype_supported
 from tinygrad.renderer.nir import NIRRenderer
 
-if getenv("TINY_BACKEND"):
+TINY_BACKEND = getenv("TINY_BACKEND")
+if TINY_BACKEND:
   import tinygrad.nn.torch # noqa: F401 # pylint: disable=unused-import
   torch.set_default_device("tiny")
 
@@ -418,7 +419,6 @@ class TestOps(unittest.TestCase):
     helper_test_op(None, lambda x: x.round(), vals=[[1.499, 1.5, 1.501, 1.0, 2.1, 0.0, -5.0, -2.499, -2.5, -2.501]], forward_only=True)
     helper_test_op(None, lambda x: x.round(), vals=[[2.5, -1.5]], forward_only=True)
 
-  @unittest.skipIf(Device.DEFAULT == "WEBGPU" and CI, "isinf check of 'nan' fails on CI software-based vulkan")
   def test_isinf(self):
     val = [float('-inf'), 0., float('inf'), float('nan'), 1.1]
     helper_test_op(None, torch.isinf, Tensor.isinf, vals=[val], forward_only=True)
@@ -640,8 +640,6 @@ class TestOps(unittest.TestCase):
     helper_test_op([(45,65), (45,65)], lambda x,y: x**y)
     helper_test_op([(45,65), (45,65)], lambda x,y: x.pow(y))
 
-  # TODO: WEBGPU NaN handling in pow operations
-  @unittest.skipIf(Device.DEFAULT == "WEBGPU", "WEBGPU NaN handling differs")
   def test_pow(self):
     helper_test_op([(45,65)], lambda x: x**0)
     helper_test_op([(45,65)], lambda x: x**1)
@@ -760,6 +758,7 @@ class TestOps(unittest.TestCase):
     data = [[1,-8,1],[32,1,6]]
     tor = torch.tensor(data, dtype=torch.int)
     ten = Tensor(data, dtype=dtypes.int32)
+    # NOTE: this breaks assigns because it's folded to 0!
     helper_test_op([], lambda: tor^tor, lambda: ten^ten, forward_only=True)
     helper_test_op([], lambda: tor^0x1337, lambda: ten^0x1337, forward_only=True)
     helper_test_op([], lambda: 0x1337^tor, lambda: 0x1337^ten, forward_only=True)
@@ -1543,7 +1542,6 @@ class TestOps(unittest.TestCase):
     helper_test_op([(3, 4, 5, 6)], lambda x: x.isclose(x + 1e-9, rtol=0.01), forward_only=True)
     helper_test_op(None, lambda x,y: x.isclose(y), vals=[[1e-7, 1e-8, 1e-9], [0.0, 0.0, 0.0]], forward_only=True)
 
-  @unittest.skipIf(Device.DEFAULT == "WEBGPU" and CI, "isinf check of 'nan' fails on CI software-based vulkan")
   def test_isclose_edge_cases(self):
     for a in [math.inf, -math.inf, math.nan, 0.0]:
       for b in [math.inf, -math.inf, math.nan, 0.0]:

test/backend/test_outerworld_call.py

@@ -1,19 +0,0 @@
-import unittest
-from tinygrad import Tensor
-
-class TestOuterCall(unittest.TestCase):
-  def test_outer_call_assign(self):
-    a = Tensor.zeros(10,10).contiguous()
-    b = Tensor.ones(10,10).contiguous()
-    Tensor.realize(a,b)
-
-    pa = a.as_param(0)
-    pb = b.as_param(1)
-    out = Tensor.call(a, b, fxn=pa.assign(pa+pb))
-    out.realize()
-
-    print(a.numpy())
-    assert (a == 1).all().item()
-
-if __name__ == '__main__':
-  unittest.main()

test/backend/test_profiler.py

@@ -1,4 +1,4 @@
-import unittest, struct, contextlib, statistics, time, gc
+import unittest, struct, contextlib, statistics, gc
 from tinygrad import Device, Tensor, dtypes, TinyJit
 from tinygrad.helpers import CI, getenv, Context, ProfileRangeEvent, cpu_profile, cpu_events, ProfilePointEvent, dedup
 from tinygrad.device import Buffer, BufferSpec, Compiled, ProfileDeviceEvent, ProfileGraphEvent
@@ -20,7 +20,7 @@ def helper_collect_profile(*devs):
   cpu_events.clear()
 
   profile_list = []
-  with Context(VIZ=1, PROFILE=1):
+  with Context(PROFILE=1):
     yield profile_list
     for dev in devs: dev.synchronize()
     for dev in devs: dev._at_profile_finalize()
@@ -170,30 +170,19 @@ class TestProfiler(unittest.TestCase):
     for (i1, d1), (i2, d2) in pairs:
       assert abs(jitter_matrix[i1][i2]) < 0.5, "jitter should be less than 0.5us"
 
-  @unittest.skip("this test is flaky")
   def test_cpu_profile(self):
     def test_fxn(err=False):
-      time.sleep(0.1)
       if err: raise Exception()
-      time.sleep(0.1)
 
     with helper_collect_profile(dev:=TestProfiler.d0) as profile:
-      with cpu_profile("test_1", dev.device):
+      with cpu_profile("test_1", dev):
         test_fxn(err=False)
       with self.assertRaises(Exception):
-        with cpu_profile("test_2", dev.device):
+        with cpu_profile("test_2", dev):
           test_fxn(err=True)
 
-    range_events = [p for p in profile if isinstance(p, ProfileRangeEvent)]
+    range_events = [p for p in profile if isinstance(p, ProfileRangeEvent) and p.device == dev]
     self.assertEqual(len(range_events), 2)
-    # record start/end time up to exit (error or success)
-    for e in range_events:
-      self.assertGreater(e.en, e.st)
-    e1, e2 = range_events
-    self.assertEqual([e1.name, e2.name], ["test_1", "test_2"])
-    # TODO: this is flaky
-    #self.assertLess(e1.st, e2.st)
-    #self.assertGreater(e1.en-e1.st, e2.en-e2.st)
 
   @unittest.skip("this test is flaky")
   @unittest.skipUnless(Device[Device.DEFAULT].graph is not None, "graph support required")

test/backend/test_renderer_failures.py

@@ -78,7 +78,9 @@ class TestCStyleFailures(unittest.TestCase):
   def test_repeat_add(self): self._test_src_strip_paren(Ops.ADD)
   def test_repeat_mul(self): self._test_src_strip_paren(Ops.MUL)
   def test_repeat_xor(self): self._test_src_strip_paren(Ops.XOR)
+  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, WGSLRenderer), "wgsl ends up with '(' * 5")
   def test_repeat_or(self): self._test_src_strip_paren(Ops.OR)
+  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, WGSLRenderer), "wgsl ends up with '(' * 5")
   def test_repeat_and(self): self._test_src_strip_paren(Ops.AND)
   def test_repeat_sub(self): self._test_src_strip_paren(Ops.SUB, should_strip_paren=False)
 

test/backend/test_schedule.py

@@ -168,13 +168,13 @@ class TestSchedule(unittest.TestCase):
     a = Tensor.full((4,), 4.0).contiguous().realize()
     b = Tensor.full((4,), 2.0).contiguous().realize()
     expr = (a*b)/b
-    run_schedule(check_schedule(expr, 0))
+    run_schedule(check_schedule(expr, 1))
     np.testing.assert_allclose(expr.numpy(), np.full((4,), 4.0))
 
   def test_div_collapse_const(self):
     a = Tensor.full((4,), 4.0).contiguous().realize()
     expr = a/a
-    run_schedule(check_schedule(expr, 0))
+    run_schedule(check_schedule(expr, 1))
     np.testing.assert_allclose(expr.numpy(), np.full((4,), 1.0))
 
   def test_div_collapse(self):
@@ -747,7 +747,7 @@ class TestSchedule(unittest.TestCase):
     p = P[0]
     p = p.pad(((1, 0), ))
     p = p.repeat([2])
-    run_schedule(check_schedule(p, 3))
+    run_schedule(check_schedule(p, 4))  # TODO: this is high
     tiny_ret = p.numpy()
 
     P = np.ones((3, 3), dtype=np.float32)
@@ -841,10 +841,9 @@ class TestSchedule(unittest.TestCase):
   def test_cast_const_view(self):
     a = Tensor.ones((4, 4), dtype=dtypes.float32)
     casted_view = a.cast(dtypes.int32)
-    run_schedule(check_schedule(casted_view, 0))
-    self.assertIsNone(casted_view.uop.base.realized)
+    run_schedule(check_schedule(casted_view, 1))
     realized_const_view = casted_view.contiguous()
-    run_schedule(check_schedule(realized_const_view, 1))
+    run_schedule(check_schedule(realized_const_view, 0))
     self.assertListEqual(realized_const_view.tolist(), [[1, 1, 1, 1], [1, 1, 1, 1], [1, 1, 1, 1], [1, 1, 1, 1]])
 
   @given(strat.sampled_from(dtypes.all), strat.sampled_from(dtypes.all))
@@ -1037,7 +1036,7 @@ class TestSchedule(unittest.TestCase):
     idx = Tensor([1,2,5,6], dtype=dtypes.int32)
     flat_base[idx] = Tensor([99,99,99,99])
     base.assign(flat_base.reshape(4, 4))
-    sched = check_schedule(base, 2)
+    sched = check_schedule(base, 6)  # TODO: this is high
     run_schedule(sched)
     expected = list(range(16))
     for i, v in zip([1,2,5,6], [99,99,99,99]): expected[i] = v
@@ -1236,8 +1235,7 @@ class TestView(unittest.TestCase):
     bv = b.pad(((0, 2),))[-2:]
     # this becomes a late a*0
     late_mul = a*bv
-    run_schedule(check_schedule(late_mul, 0))
-    # NOTE: no longer checked
+    run_schedule(check_schedule(late_mul, 2))
     # the arange doesn't realize
     #self.assertIsNone(b.uop.base.realized)
     # mul doesn't realize
@@ -1254,7 +1252,7 @@ class TestView(unittest.TestCase):
     bv = b.pad(((0, 2),))[-2:]
     late_mul = a*bv
     other_child = b+2
-    s = check_schedule([late_mul, other_child], 2)
+    s = check_schedule([late_mul, other_child], 3)
     # the arange becomes a BUFFER
     self.assertIs(b.uop.base.op, Ops.BUFFER)
     # NOTE: no longer checked
@@ -1267,7 +1265,7 @@ class TestView(unittest.TestCase):
 class TestCopyFolding(unittest.TestCase):
   def test_const_copy_is_free(self):
     b = Tensor(1).to("CPU") * 4
-    run_schedule(check_schedule(b, 0, filter_sink=False))
+    run_schedule(check_schedule(b, 1, filter_sink=False))
     assert b.item() == 4
 
   def test_one_hot_with_copy(self):
@@ -1277,14 +1275,14 @@ class TestCopyFolding(unittest.TestCase):
 
   def test_const_copy_multi(self):
     x = Tensor.ones(1, device="CPU").to_(["CPU", "CPU:1"]) * 2
-    run_schedule(check_schedule(x, 0, filter_sink=False))
+    run_schedule(check_schedule(x, 2, filter_sink=False))
     self.assertEqual(x.item(), 2.0)
 
   def test_late_const_copy_folding(self):
     a = Tensor.arange(3).realize()
     zeros = Tensor.zeros(3).realize()
     b = (a*zeros).to("CPU") + 1
-    run_schedule(check_schedule(b, 0, filter_sink=False))
+    run_schedule(check_schedule(b, 1, filter_sink=False))
     self.assertListEqual(b.tolist(), [1, 1, 1])
     self.assertEqual(b.device, "CPU")
 
@@ -1324,7 +1322,7 @@ class TestCopyFolding(unittest.TestCase):
     a = Tensor.ones(4, 4).contiguous().realize()
     # use copy_to_device to bypass Tensor.to() shortcircuit and force a real same-device COPY in the graph
     a.assign(Tensor(a.uop.copy_to_device(a.device), a.device))
-    run_schedule(check_schedule(a, 0, filter_sink=False))
+    run_schedule(check_schedule(a, 2, filter_sink=False))
     self.assertListEqual(a.tolist(), [[1.]*4]*4)
 
   def test_clone(self):

test/backend/test_symbolic_jit.py

@@ -80,7 +80,7 @@ class TestSymbolicJit(unittest.TestCase):
       symbolic = jf(q, k[:, :vi], v[:, :vi])[:2, :4, :1, :8].numpy()
       expected = f(q, k[:, :i], v[:, :i]).numpy()
       np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
-    assert_jit_cache_len(jf, 4)
+    assert_jit_cache_len(jf, 5)
 
   def test_cat_dim0(self):
     def f(a, b): return a.cat(b, dim=0).realize()

test/backend/test_transcendental.py

@@ -84,7 +84,6 @@ class TestFromFuzzer(unittest.TestCase):
     _test_value(np.pi * 2, unit=1.5)
 
   @given(strat.sampled_from(dtypes_float))
-  @unittest.skipIf(Device.DEFAULT == "WEBGPU" and CI, "Nan location mismatch on Vulkan, Metal works")
   def test_log2(self, dtype):
     if not is_dtype_supported(dtype): return
     if dtype == dtypes.float64:

test/backend/test_uops.py

@@ -113,6 +113,12 @@ class TestFloatUOps(TestUOps):
   def test_max(self): self._test_bop_fxn(Ops.MAX, lambda a,b: max(a,b))
   def test_cmplt(self): self._test_bop_fxn(Ops.CMPLT, lambda a,b: a<b)
   def test_cmpne(self): self._test_bop_fxn(Ops.CMPNE, lambda a,b: a!=b)
+  @unittest.skipIf(Device.DEFAULT == "WEBGPU", "WEBGPU doesn't support NaN comparison correctly")
+  def test_cmpne_nan(self):  # NaN != x for any x (IEEE 754)
+    for a, b in [(math.nan, 1.0), (1.0, math.nan), (math.nan, math.nan)]:
+      self.assertTrue(_test_single_value(
+        [dtypes.as_const(a, dtypes.float32), dtypes.as_const(b, dtypes.float32)],
+        Ops.CMPNE, (dtypes.float32, dtypes.float32)))
   # MOD isn't tested on floats
 
   def test_where(self):

test/device/test_hcq.py

@@ -76,7 +76,7 @@ class TestHCQ(unittest.TestCase):
         TestHCQ.d0.timeline_signal.wait(TestHCQ.d0.timeline_value)
         TestHCQ.d0.timeline_value += 1
 
-  @unittest.skipIf(Device.DEFAULT in {"CPU"}, "Can't handle async update on CPU device")
+  @unittest.skipIf(Device.DEFAULT in {"CPU"} or getenv("AMD_IFACE", "") == "PCI", "Can't handle async update on CPU/MOCKAM device")
   def test_wait_late_set(self):
     for queue_type in [TestHCQ.d0.hw_compute_queue_t, TestHCQ.d0.hw_copy_queue_t]:
       if queue_type is None: continue
@@ -538,7 +538,7 @@ class TestHCQ(unittest.TestCase):
 
       np.testing.assert_equal(cpu_buffer.numpy(), local_buf.numpy(), "failed")
 
-  @unittest.skipUnless(MOCKGPU, "Emulate this on MOCKGPU to check the path in CI")
+  @unittest.skipUnless(MOCKGPU and getenv("AMD_IFACE", "") != "PCI", "Emulate this on MOCKGPU to check the path in CI")
   def test_on_device_hang(self):
     if not hasattr(self.d0, 'on_device_hang'): self.skipTest("device does not have on_device_hang")
 

test/device/test_metal.py

@@ -1,5 +1,5 @@
 import unittest
-from tinygrad.device import CompileError, Device
+from tinygrad.device import CompileError, Device, BufferSpec
 if Device.DEFAULT=="METAL":
   from tinygrad.runtime.ops_metal import MetalDevice, MetalCompiler, MetalProgram
 @unittest.skipIf(Device.DEFAULT!="METAL", "Metal support required")
@@ -48,4 +48,14 @@ kernel void r_5(device int* data0, const device int* data1, uint3 gid [[threadgr
 """)
     with self.assertRaises(RuntimeError):
       compiled = compiled[:40] # corrupt the compiled program
-      MetalProgram(device, "r_5", compiled)
+      MetalProgram(device, "r_5", compiled)
+
+  def test_free(self):
+    size = 2**16
+    device = Device['METAL']
+    before = device.sysdevice.currentAllocatedSize()
+
+    buf = device.allocator.alloc(size, BufferSpec(nolru=True))
+    self.assertEqual(curr:=device.sysdevice.currentAllocatedSize(), before+size, msg=f"{curr=} - {before=}")
+    device.allocator.free(buf, buf.size, BufferSpec(nolru=True))
+    self.assertEqual(curr:=device.sysdevice.currentAllocatedSize(), before, msg=f"{curr=} - {before=}")

test/external/external_test_hive_reset.py

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

test/external/process_replay/process_replay.py

@@ -1,6 +1,6 @@
 #!/usr/bin/env python3
 # compare kernels created by HEAD against master
-import os, multiprocessing, logging, pickle, sqlite3, difflib, warnings, itertools, functools, base64, codecs
+import os, multiprocessing, logging, pickle, sqlite3, difflib, warnings, functools, base64, codecs
 from dataclasses import replace
 from typing import Callable, Any
 
@@ -8,7 +8,6 @@ ASSERT_DIFF = int((flag:="[pr]") in os.getenv("COMMIT_MESSAGE", flag) or flag in
 if not int(os.getenv("ASSERT_PROCESS_REPLAY", "1")): ASSERT_DIFF = 0
 
 try:
-  from tinygrad.schedule.rangeify import get_rangeify_map
   from tinygrad.renderer import Renderer, ProgramSpec
   from tinygrad.engine.realize import get_program
   from tinygrad.uop.ops import UOp, Ops, KernelInfo
@@ -43,14 +42,6 @@ class ProcessReplayWarning(Warning): pass
 
 # *** replay the function and convert return values to string
 
-def replay_get_rangeify_map(ret:dict[UOp, UOp], big_sink:UOp) -> tuple[str, str, tuple[Any, ...]]:
-  UOp.unique_num = itertools.count(max([u.arg for u in big_sink.toposort() if u.op is Ops.UNIQUE], default=0)+1)
-  new_sink = big_sink.substitute(get_rangeify_map(big_sink))
-  def to_str(ret:UOp) -> str:
-    asts = [repr(u.arg.ast) for u in ret.toposort() if u.op is Ops.CALL]
-    return "\n".join([f"{len(asts)} kernels", *asts])
-  return to_str(new_sink), to_str(big_sink.substitute(ret)), (big_sink,)
-
 def replay_get_program(p:ProgramSpec, ast:UOp, renderer:Renderer, opts:list[Opt]|None=None) -> tuple[str, str, tuple[Any, ...]]:
   # the ast.arg is non None if we are inside of search.py
   sink_arg = ast.arg or KernelInfo()
@@ -68,8 +59,6 @@ def replay_get_program(p:ProgramSpec, ast:UOp, renderer:Renderer, opts:list[Opt]
 
 replayers: dict[str, Callable[..., tuple[str, str, tuple[Any, ...]]]] = {}
 replayers["get_program"] = replay_get_program
-# disable this for speed, does it ever find things?
-#replayers["get_rangeify_map"] = replay_get_rangeify_map
 
 # *** run replayers on captured rows and print diffs
 

test/helpers.py

@@ -39,7 +39,7 @@ def assert_jit_cache_len(fxn, expected_len):
     assert len(fxn.jit_cache) == 1, len(fxn.jit_cache)
     # until we have a better way of typing the prg in ExecItem
     assert type(fxn.jit_cache[0].prg).__name__.endswith('Graph')
-    assert len(fxn.jit_cache[0].prg.jit_cache) == expected_len
+    assert len(fxn.jit_cache[0].prg.jit_cache) == expected_len, f"expected {expected_len}, got {len(fxn.jit_cache[0].prg.jit_cache)}"
 
 def rand_for_dtype(dt:DType, size:int, allow_subnormal=True):
   if dtypes.is_unsigned(dt):

test/mockgpu/am/amdriver.py

@@ -0,0 +1,127 @@
+from __future__ import annotations
+import mmap, functools
+from tinygrad.runtime.autogen import libc
+from test.mockgpu.driver import VirtDriver, VirtFileDesc, TextFileDesc, DirFileDesc, VirtFile
+from test.mockgpu.am.amgpu import MockAMGPU, VRAM_SIZE
+
+DOORBELL_SIZE = 0x2000
+MMIO_SIZE = 2 << 20
+PCIBUS = "mock:am:0"
+
+_empty_bar = "0x0000000000000000 0x0000000000000000 0x0000000000000000"
+_resource_lines = [
+  f"0x0000000000000000 0x{VRAM_SIZE-1:016x} 0x0000000000000000", _empty_bar,
+  f"0x0000000000000000 0x{DOORBELL_SIZE-1:016x} 0x0000000000000000", _empty_bar, _empty_bar,
+  f"0x0000000000000000 0x{MMIO_SIZE-1:016x} 0x0000000000000000", _empty_bar,
+]
+
+class PagemapFileDesc(VirtFileDesc):
+  def __init__(self, fd, gpu):
+    super().__init__(fd)
+    self.gpu = gpu
+  def seek(self, offset): self.off = offset
+  def read_contents(self, size=None):
+    entries = bytearray()
+    for i in range((size or 8) // 8):
+      vaddr = ((self.off // 8) + i) * 0x1000
+      paddr = self.gpu._next_sysmem_paddr
+      self.gpu._next_sysmem_paddr += 0x1000
+      self.gpu._sysmem_map[paddr] = vaddr
+      entries += ((1 << 63) | (paddr // 0x1000)).to_bytes(8, 'little')
+    self.off += len(entries)
+    return bytes(entries)
+
+class PCIBarFileDesc(VirtFileDesc):
+  def __init__(self, fd, memfd, driver=None):
+    super().__init__(fd)
+    self.memfd, self.driver = memfd, driver
+  def mmap(self, start, sz, prot, flags, fd, off):
+    addr = libc.mmap(start, sz, prot, flags, self.memfd, off)
+    if self.driver is not None:
+      self.driver.track_address(addr, addr + sz, lambda mv, idx: None, lambda mv, idx: self.driver._emulate_execute())
+    return addr
+
+class PCIMMIOBarFileDesc(VirtFileDesc):
+  def __init__(self, fd, bar5_addr):
+    super().__init__(fd)
+    self.bar5_addr = bar5_addr
+  def mmap(self, start, sz, prot, flags, fd, off): return self.bar5_addr + off
+
+class PCIConfigFileDesc(VirtFileDesc):
+  def __init__(self, fd):
+    super().__init__(fd)
+    self.data = bytearray(256)
+  def read_contents(self, size=None): return bytes(self.data[self.off:self.off + (size or len(self.data) - self.off)])
+  def write_contents(self, content): self.data[self.off:self.off + len(content)] = content
+  def seek(self, offset): self.off = offset
+
+class PCIEnableFileDesc(VirtFileDesc):
+  def __init__(self, fd): super().__init__(fd)
+  def read_contents(self, size=None): return "1\n"
+  def write_contents(self, content): pass
+
+class AMDriver(VirtDriver):
+  def __init__(self):
+    super().__init__()
+    self.gpus:dict[int, MockAMGPU] = {}
+    self._executing = False
+    self.gpu = MockAMGPU(0)
+    self.gpus[0] = self.gpu
+    self.next_fd = 1 << 30
+
+    self._bar5_addr = libc.mmap(0, MMIO_SIZE, mmap.PROT_READ | mmap.PROT_WRITE, mmap.MAP_SHARED | mmap.MAP_ANONYMOUS, -1, 0)
+    mmio = self.gpu.mmio
+    self.track_address(self._bar5_addr, self._bar5_addr + MMIO_SIZE,
+      lambda mv, idx: _bar5_sync_read(mv, idx, mmio), lambda mv, idx: _bar5_sync_write(mv, idx, mmio))
+
+    p = f"/sys/bus/pci/devices/{PCIBUS}"
+    self.tracked_files += [
+      VirtFile("/proc/sys/vm/compact_unevictable_allowed", functools.partial(TextFileDesc, text="0\n")),
+      VirtFile("/proc/self/pagemap", functools.partial(PagemapFileDesc, gpu=self.gpu)),
+      VirtFile("/sys/bus/pci/devices", functools.partial(DirFileDesc, child_names=[PCIBUS])),
+      VirtFile(f"{p}/vendor", functools.partial(TextFileDesc, text="0x1002\n")),
+      VirtFile(f"{p}/device", functools.partial(TextFileDesc, text="0x74a1\n")),
+      VirtFile(f"{p}/enable", PCIEnableFileDesc),
+      VirtFile(f"{p}/config", PCIConfigFileDesc),
+      VirtFile(f"{p}/resource", functools.partial(TextFileDesc, text="\n".join(_resource_lines) + "\n")),
+      VirtFile(f"{p}/resource0", functools.partial(PCIBarFileDesc, memfd=self.gpu.vram_fd)),
+      VirtFile(f"{p}/resource2", functools.partial(PCIBarFileDesc, memfd=self.gpu.doorbell_fd, driver=self)),
+      VirtFile(f"{p}/resource5", functools.partial(PCIMMIOBarFileDesc, bar5_addr=self._bar5_addr)),
+    ]
+
+  def _alloc_fd(self):
+    fd = self.next_fd
+    self.next_fd += 1
+    return fd
+
+  def open(self, name, flags, mode, virtfile): return virtfile.fdcls(self._alloc_fd())
+
+  def _emulate_execute(self):
+    if self._executing: return
+    self._executing = True
+    try:
+      any_progress = True
+      while any_progress:
+        any_progress = False
+        for gpu in self.gpus.values():
+          for q in gpu.queues:
+            if q.executing: any_progress |= q.execute() > 0
+    finally:
+      self._executing = False
+
+def _bar5_sync_read(mv, idx, mmio):
+  if isinstance(idx, slice):
+    for i in range(idx.start or 0, idx.stop or len(mv), idx.step or 1): mv[i] = mmio[i]
+  else: mv[idx] = mmio[idx]
+
+def _bar5_sync_write(mv, idx, mmio):
+  if isinstance(idx, slice):
+    for i in range(idx.start or 0, idx.stop or len(mv), idx.step or 1): mmio[i] = mv[i]
+  else: mmio[idx] = mv[idx]
+
+class AMUSBDriver(AMDriver):
+  def __init__(self):
+    import test.mockgpu.usb as _musb
+    super().__init__()
+    self.state = _musb.MockASM24State(self.gpu, self, VRAM_SIZE, DOORBELL_SIZE, MMIO_SIZE)
+    _musb._mock_usb_state = self.state

test/mockgpu/am/amgpu.py

@@ -0,0 +1,314 @@
+# mypy: ignore-errors
+from __future__ import annotations
+import ctypes, ctypes.util, struct, functools, os, mmap
+from tinygrad.runtime.autogen.am import am
+from tinygrad.runtime.support.amd import AMDReg, import_asic_regs
+from test.mockgpu.amd.amdgpu import AMDGPU
+
+libc = ctypes.CDLL(ctypes.util.find_library("c"))
+libc.mmap.argtypes = [ctypes.c_void_p, ctypes.c_size_t, ctypes.c_int, ctypes.c_int, ctypes.c_int, ctypes.c_long]
+libc.mmap.restype = ctypes.c_void_p
+
+VRAM_SIZE = 512 << 20
+
+IP_VERSIONS = {
+  am.GC_HWIP: (12, 0, 0), am.SDMA0_HWIP: (7, 0, 0), am.MMHUB_HWIP: (4, 1, 0), am.NBIO_HWIP: (6, 3, 1),
+  am.MP0_HWIP: (14, 0, 2), am.MP1_HWIP: (14, 0, 2), am.HDP_HWIP: (7, 0, 0), am.OSSSYS_HWIP: (7, 0, 0),
+}
+
+def _pad(t, n=10): return t + (0,) * (n - len(t))
+IP_BASES = {
+  am.GC_HWIP:     _pad((0x00001260, 0x0000A000, 0x0001C000, 0x02402C00)),
+  am.SDMA0_HWIP:  _pad((0x00001260, 0x0000A000, 0x0001C000, 0x02402C00)),
+  am.MMHUB_HWIP:  _pad((0x0001A000, 0x02408800)),
+  am.NBIO_HWIP:   _pad((0x00000000, 0x00000014, 0x00000D20, 0x00010400, 0x0241B000, 0x04040000)),
+  am.MP0_HWIP:    _pad((0x00016000, 0x00DC0000, 0x00E00000, 0x00E40000, 0x0243FC00)),
+  am.MP1_HWIP:    _pad((0x00016000, 0x00DC0000, 0x00E00000, 0x00E40000, 0x0243FC00)),
+  am.HDP_HWIP:    _pad((0x00000F20, 0x0240A400)),
+  am.OSSSYS_HWIP: _pad((0x000010A0, 0x0240A000)),
+}
+
+IP_HWIDS = {hwip: am.hw_id_map[hwip] for hwip in IP_VERSIONS}
+
+GC_INFO = dict(gc_num_se=2, gc_num_cu_per_sh=8, gc_num_sh_per_se=2, gc_num_rb_per_se=4,
+               gc_num_tccs=8, gc_wave_size=32, gc_max_waves_per_simd=16, gc_max_scratch_slots_per_cu=32, gc_lds_size=64)
+
+def _build_ip_regs(prefix, hwip) -> dict[str, AMDReg]:
+  try: return import_asic_regs(prefix, IP_VERSIONS[hwip], cls=functools.partial(AMDReg, bases={0: IP_BASES[hwip]}))
+  except Exception: return {}
+
+class MockMMU:
+  def __init__(self, gpu:MockAMGPU):
+    self.gpu = gpu
+    self.tlb: dict[int, tuple[int, int, bool]] = {}
+
+  def invalidate(self, pt_base:int, va_base:int):
+    new_tlb: dict[int, tuple[int, int, bool]] = {}
+    self._walk(pt_base, 0, 0, new_tlb, va_base)
+    for va, (pa, sz, is_sys) in new_tlb.items():
+      old = self.tlb.get(va)
+      if not is_sys and (old is None or old[0] != pa): self.gpu.map_vram_at(va, pa, sz)
+      if old is None: self.gpu.map_range(va, sz)
+    self.tlb = new_tlb
+
+  def _walk(self, pt_paddr:int, level:int, va_acc:int, out:dict, va_base:int):
+    shift = [39, 30, 21, 12][level]
+    for i in range(512):
+      pte = struct.unpack_from('<Q', self.gpu.vram, pt_paddr + i * 8)[0]
+      if not (pte & am.AMDGPU_PTE_VALID): continue
+      va, pa = va_acc | (i << shift), pte & 0x0000FFFFFFFFF000
+      if level == 3 or (pte & am.AMDGPU_PDE_PTE_GFX12):
+        out[va_base + va] = (pa, 1 << shift, bool(pte & am.AMDGPU_PTE_SYSTEM))
+      else:
+        self._walk(pa, level + 1, va, out, va_base)
+
+  def paddr_to_host(self, paddr:int) -> int:
+    page, off = paddr & ~0xFFF, paddr & 0xFFF
+    if page in self.gpu._sysmem_map: return self.gpu._sysmem_map[page] + off
+    if paddr < VRAM_SIZE: return self.gpu.vram_addr + paddr
+    raise ValueError(f"paddr {paddr:#x} not found in sysmem_map or VRAM")
+
+  def addr_to_host(self, addr:int) -> int:
+    gmc = self.gpu.mmio.gmc
+    sys_lo = self.gpu.mmio.regs.get(gmc.reg('regMMMC_VM_SYSTEM_APERTURE_LOW_ADDR') or 0, 0) << 18
+    sys_hi = self.gpu.mmio.regs.get(gmc.reg('regMMMC_VM_SYSTEM_APERTURE_HIGH_ADDR') or 0, 0) << 18
+    if sys_lo <= addr < sys_hi: return self.paddr_to_host(addr - self.gpu.mc_base)
+    for tva, (pa, sz, is_sys) in self.tlb.items():
+      if tva <= addr < tva + sz:
+        paddr = pa + (addr - tva)
+        if not is_sys: return self.gpu.vram_addr + paddr
+        return self.paddr_to_host(paddr)
+    raise ValueError(f"addr {addr:#x} not mapped (sys_aperture=[{sys_lo:#x}, {sys_hi:#x}])")
+
+class MockIPBlock:
+  def __init__(self, gpu:MockAMGPU, mmio:MockMMIOInterface, regs:dict[str, AMDReg]):
+    self.gpu, self.mmio, self._regs = gpu, mmio, regs
+    self._n2a = {n: r.addr[0] for n, r in regs.items()}
+    self._a2n = {a: n for n, a in self._n2a.items()}
+    self.addrs = set(self._n2a.values())
+  def reg(self, name) -> int|None: return self._n2a.get(name)
+  def decode(self, name) -> dict: return self._regs[name].decode(self.mmio.regs.get(self._n2a[name], 0))
+  def read(self, reg:int) -> int: return self.mmio.regs.get(reg, 0)
+  def write(self, reg:int, val:int): self.mmio.regs[reg] = val
+  def _read_pair(self, pair) -> int:
+    if pair[0] is None: return 0
+    return self.mmio.regs.get(pair[0], 0) | (self.mmio.regs.get(pair[1], 0) << 32)
+
+class MockPSP(MockIPBlock):
+  def __init__(self, gpu, mmio):
+    super().__init__(gpu, mmio, _build_ip_regs('mp', am.MP0_HWIP))
+    self._sos_alive, self._ring_wptr = False, 0
+    pref = "regMPASP_SMN_C2PMSG" if IP_VERSIONS[am.MP0_HWIP] >= (14,0,0) else "regMP0_SMN_C2PMSG"
+    def r(n): return self.reg(f"{pref}_{n}")
+    self._c2pmsg_35, self._c2pmsg_64, self._c2pmsg_67 = r(35), r(64), r(67)
+    self._c2pmsg_69, self._c2pmsg_70, self._c2pmsg_81 = r(69), r(70), r(81)
+
+  def read(self, reg:int) -> int:
+    if reg == self._c2pmsg_35: return 0x80000000
+    if reg == self._c2pmsg_81: return 0x1 if self._sos_alive else 0x0
+    if reg == self._c2pmsg_64: return 0x80000000 if self._sos_alive else 0x0
+    if reg == self._c2pmsg_67: return self._ring_wptr
+    return super().read(reg)
+
+  def write(self, reg:int, val:int):
+    super().write(reg, val)
+    if reg == self._c2pmsg_35 and val == am.PSP_BL__LOAD_SOSDRV: self._sos_alive = True
+    if reg == self._c2pmsg_67: self._ring_submit(val)
+
+  def _ring_submit(self, new_wptr:int):
+    old_wptr = self._ring_wptr
+    self._ring_wptr = new_wptr
+    lo, hi = self._c2pmsg_69, self._c2pmsg_70
+    if lo is None or hi is None: return
+    ring_mc = self.mmio.regs.get(lo, 0) | (self.mmio.regs.get(hi, 0) << 32)
+    ring_paddr = ring_mc - self.gpu.mc_base
+    frame_off = ring_paddr + old_wptr * 4
+    frame = am.struct_psp_gfx_rb_frame.from_buffer_copy(bytes(self.gpu.vram[frame_off:frame_off + ctypes.sizeof(am.struct_psp_gfx_rb_frame)]))
+    fence_paddr = ((frame.fence_addr_hi << 32) | frame.fence_addr_lo) - self.gpu.mc_base
+    if 0 <= fence_paddr < len(self.gpu.vram):
+      struct.pack_into('<I', self.gpu.vram, fence_paddr, frame.fence_value)
+    cmd_paddr = ((frame.cmd_buf_addr_hi << 32) | frame.cmd_buf_addr_lo) - self.gpu.mc_base
+    if 0 <= cmd_paddr < len(self.gpu.vram):
+      struct.pack_into('<I', self.gpu.vram, cmd_paddr + 864, 0)
+
+class MockSMU(MockIPBlock):
+  def __init__(self, gpu, mmio):
+    try: regs = import_asic_regs('mp', (11, 0), cls=functools.partial(AMDReg, bases={0: IP_BASES[am.MP1_HWIP]}))
+    except Exception: regs = {}
+    super().__init__(gpu, mmio, regs)
+    self._msg_pending = False
+    def r(n): return self.reg(f"mmMP1_SMN_C2PMSG_{n}")
+    self._c2pmsg_53, self._c2pmsg_54, self._c2pmsg_66 = r(53), r(54), r(66)
+    self._c2pmsg_75, self._c2pmsg_82, self._c2pmsg_90 = r(75), r(82), r(90)
+
+  def read(self, reg:int) -> int:
+    if reg == self._c2pmsg_90 or reg == self._c2pmsg_54: return 0x1 if self._msg_pending else super().read(reg)
+    if reg == self._c2pmsg_82: return self.mmio.regs.get(reg, 3)
+    return super().read(reg)
+
+  def write(self, reg:int, val:int):
+    super().write(reg, val)
+    if reg == self._c2pmsg_66 or reg == self._c2pmsg_75: self._msg_pending = True
+    if (reg == self._c2pmsg_90 or reg == self._c2pmsg_54) and val == 0: self._msg_pending = False
+
+class MockSDMA(MockIPBlock):
+  def __init__(self, gpu, mmio):
+    all_gc = _build_ip_regs('gc', am.GC_HWIP)
+    super().__init__(gpu, mmio, {n: r for n, r in all_gc.items() if 'SDMA' in n})
+
+  def write(self, reg:int, val:int):
+    super().write(reg, val)
+    name = self._a2n.get(reg, '')
+    if name.endswith('_RB_CNTL') and self._regs[name].decode(val).get('rb_enable', 0):
+      self._activate_queue(name.rsplit('_RB_CNTL', 1)[0])
+
+  def _activate_queue(self, prefix:str):
+    ring_addr = self._read_pair((self.reg(f'{prefix}_RB_BASE'), self.reg(f'{prefix}_RB_BASE_HI'))) << 8
+    rptr_addr = self._read_pair((self.reg(f'{prefix}_RB_RPTR_ADDR_LO'), self.reg(f'{prefix}_RB_RPTR_ADDR_HI')))
+    wptr_addr = self._read_pair((self.reg(f'{prefix}_RB_WPTR_POLL_ADDR_LO'), self.reg(f'{prefix}_RB_WPTR_POLL_ADDR_HI')))
+    rb_size = self.decode(f'{prefix}_RB_CNTL')['rb_size']
+    self.gpu.add_sdma_queue(self.gpu.mmu.addr_to_host(ring_addr), 4 << rb_size,
+                            self.gpu.mmu.addr_to_host(rptr_addr), self.gpu.mmu.addr_to_host(wptr_addr))
+
+class MockGFX(MockIPBlock):
+  def __init__(self, gpu, mmio):
+    super().__init__(gpu, mmio, _build_ip_regs('gc', am.GC_HWIP))
+    self._pt_base = (self.reg('regGCVM_CONTEXT0_PAGE_TABLE_BASE_ADDR_LO32'), self.reg('regGCVM_CONTEXT0_PAGE_TABLE_BASE_ADDR_HI32'))
+    self._pt_start = (self.reg('regGCVM_CONTEXT0_PAGE_TABLE_START_ADDR_LO32'), self.reg('regGCVM_CONTEXT0_PAGE_TABLE_START_ADDR_HI32'))
+    self._gc_inv_ack = self.reg('regGCVM_INVALIDATE_ENG17_ACK')
+    self._gc_inv_req = self.reg('regGCVM_INVALIDATE_ENG17_REQ')
+    self._hqd_active = self.reg('regCP_HQD_ACTIVE')
+
+  def read(self, reg:int) -> int:
+    if reg == self.reg('regCP_STAT') or reg == self.reg('regRLC_SAFE_MODE'): return 0
+    if reg == self.reg('regRLC_RLCS_BOOTLOAD_STATUS'): return 0x2
+    if reg == self._gc_inv_ack: return 0x1
+    return super().read(reg)
+
+  def write(self, reg:int, val:int):
+    super().write(reg, val)
+    if reg == self.reg('regCP_HQD_DEQUEUE_REQUEST'):
+      if self._hqd_active is not None: self.mmio.regs[self._hqd_active] = 0
+    if reg == self._hqd_active and val == 1: self._activate_pm4_queue()
+    if reg == self._gc_inv_req: self.gpu.mmu.invalidate(self.get_pt_base(), self.get_va_base())
+
+  def _activate_pm4_queue(self):
+    ring_addr = self._read_pair((self.reg('regCP_HQD_PQ_BASE'), self.reg('regCP_HQD_PQ_BASE_HI'))) << 8
+    rptr_addr = self._read_pair((self.reg('regCP_HQD_PQ_RPTR_REPORT_ADDR'), self.reg('regCP_HQD_PQ_RPTR_REPORT_ADDR_HI')))
+    wptr_addr = self._read_pair((self.reg('regCP_HQD_PQ_WPTR_POLL_ADDR'), self.reg('regCP_HQD_PQ_WPTR_POLL_ADDR_HI')))
+    queue_size = self.decode('regCP_HQD_PQ_CONTROL')['queue_size']
+    self.gpu.add_pm4_queue(self.gpu.mmu.addr_to_host(ring_addr), 4 << (queue_size + 1),
+                           self.gpu.mmu.addr_to_host(rptr_addr), self.gpu.mmu.addr_to_host(wptr_addr))
+
+  def get_pt_base(self) -> int: return self._read_pair(self._pt_base) & 0x0000FFFFFFFFF000
+  def get_va_base(self) -> int: return self._read_pair(self._pt_start) << 12
+
+class MockGMC(MockIPBlock):
+  def __init__(self, gpu, mmio, gfx:MockGFX):
+    super().__init__(gpu, mmio, _build_ip_regs('mmhub', am.MMHUB_HWIP))
+    self._gfx = gfx
+    self._inv_ack = self.reg('regMMVM_INVALIDATE_ENG17_ACK')
+    self._inv_sem = self.reg('regMMVM_INVALIDATE_ENG17_SEM')
+    self._inv_req = self.reg('regMMVM_INVALIDATE_ENG17_REQ')
+    self._fb_loc_top = self.reg('regMMMC_VM_FB_LOCATION_TOP')
+
+  def read(self, reg:int) -> int:
+    if reg == self._inv_ack or reg == self._inv_sem: return 0x1
+    if reg == self._fb_loc_top: return VRAM_SIZE >> 24
+    return super().read(reg)
+
+  def write(self, reg:int, val:int):
+    super().write(reg, val)
+    if reg == self._inv_req: self.gpu.mmu.invalidate(self._gfx.get_pt_base(), self._gfx.get_va_base())
+
+class MockNBIO(MockIPBlock):
+  def __init__(self, gpu, mmio):
+    regs = _build_ip_regs('nbif', am.NBIO_HWIP)
+    regs.update(_build_ip_regs('hdp', am.HDP_HWIP))
+    super().__init__(gpu, mmio, regs)
+    self._remap_hdp = self.reg('regBIF_BX0_REMAP_HDP_MEM_FLUSH_CNTL')
+    self._hdp_flush = self.reg('regHDP_MEM_FLUSH_CNTL')
+
+  def read(self, reg:int) -> int:
+    if reg == self._remap_hdp and self._hdp_flush is not None: return self._hdp_flush * 4
+    return super().read(reg)
+
+class MockMMIOInterface:
+  def __init__(self, gpu:MockAMGPU):
+    self.gpu = gpu
+    self.regs: dict[int, int] = {}
+    gfx = MockGFX(gpu, self)
+    self.gmc = MockGMC(gpu, self, gfx)
+    self.blocks = [MockPSP(gpu, self), MockSMU(gpu, self), MockSDMA(gpu, self), gfx, self.gmc, MockNBIO(gpu, self)]
+    self._addr_block: dict[int, MockIPBlock] = {}
+    for block in self.blocks:
+      for addr in block.addrs: self._addr_block.setdefault(addr, block)
+
+  def __getitem__(self, index:int|slice) -> int|list[int]:
+    if isinstance(index, slice): return [self[i] for i in range(index.start or 0, index.stop or 0, index.step or 1)]  # type: ignore[misc]
+    if index == 0xde3: return VRAM_SIZE >> 20
+    if block := self._addr_block.get(index): return block.read(index)
+    return self.regs.get(index, 0)
+
+  def __setitem__(self, index:int|slice, val:int|list[int]|tuple[int, ...]):
+    if isinstance(index, slice):
+      vals = val if isinstance(val, (list, tuple)) else [val] * ((index.stop - index.start) // (index.step or 1))  # type: ignore[operator]
+      for i, v in zip(range(index.start or 0, index.stop or 0, index.step or 1), vals): self[i] = v
+      return
+    assert isinstance(val, int)
+    self.regs[index] = val
+    if block := self._addr_block.get(index): block.write(index, val)
+
+  def __len__(self): return 0x10000000
+
+class MockAMGPU(AMDGPU):
+  def __init__(self, gpuid:int=0):
+    super().__init__(gpuid)
+    self.vram_fd = os.memfd_create("vram")
+    os.ftruncate(self.vram_fd, VRAM_SIZE)
+    self.vram_addr = libc.mmap(0, VRAM_SIZE, mmap.PROT_READ | mmap.PROT_WRITE, mmap.MAP_SHARED, self.vram_fd, 0)
+    self.vram = (ctypes.c_ubyte * VRAM_SIZE).from_address(self.vram_addr)
+    self.doorbell_fd = os.memfd_create("doorbell")
+    os.ftruncate(self.doorbell_fd, 0x2000)
+    self.arch = "rdna4"
+    self._sysmem_map:dict[int,int] = {}
+    self._next_sysmem_paddr = 0x100000000
+    self.mmu = MockMMU(self)
+    self.mmio = MockMMIOInterface(self)
+    self._preboot()
+
+  def translate_addr(self, addr:int) -> int: return self.mmu.addr_to_host(addr)
+
+  def map_vram_at(self, va:int, paddr:int, size:int):
+    libc.mmap(va, size, mmap.PROT_READ | mmap.PROT_WRITE, mmap.MAP_SHARED | 0x10, self.vram_fd, paddr)
+
+  def _preboot(self):
+    ip_data = bytearray()
+    for hwip, (major, minor, rev) in IP_VERSIONS.items():
+      ip = am.struct_ip_v4(hw_id=IP_HWIDS[hwip], num_base_address=len(IP_BASES[hwip]), major=major, minor=minor, revision=rev)
+      ip_data += bytes(ip) + b'\x00'
+      for b in IP_BASES[hwip]: ip_data += struct.pack('<I', b)
+
+    dhdr = am.struct_die_header(num_ips=len(IP_VERSIONS))
+    ihdr = am.struct_ip_discovery_header(signature=am.DISCOVERY_TABLE_SIGNATURE, version=4, num_dies=1)
+    ip_disc_off = ctypes.sizeof(am.struct_binary_header)
+    ihdr.die_info[0].die_offset = ip_disc_off + ctypes.sizeof(am.struct_ip_discovery_header)
+
+    gc = am.struct_gc_info_v2_1()
+    gc.header.table_id, gc.header.version_major, gc.header.version_minor = am.GC, 2, 1
+    gc.header.size = ctypes.sizeof(am.struct_gc_info_v2_1)
+    for field, val in GC_INFO.items(): setattr(gc, field, val)
+
+    gc_off = ip_disc_off + ctypes.sizeof(am.struct_ip_discovery_header) + ctypes.sizeof(am.struct_die_header) + len(ip_data)
+    bhdr = am.struct_binary_header(binary_signature=am.BINARY_SIGNATURE)
+    bhdr.table_list[am.IP_DISCOVERY].offset = ip_disc_off
+    bhdr.table_list[am.GC].offset = gc_off
+
+    tbl = bytes(bhdr) + bytes(ihdr) + bytes(dhdr) + ip_data + bytes(gc)
+    tbl_offset = VRAM_SIZE - (64 << 10)
+    self.vram[tbl_offset:tbl_offset + len(tbl)] = list(tbl)
+
+  @property
+  def mc_base(self) -> int:
+    fb_loc_base = self.mmio.gmc.reg('regMMMC_VM_FB_LOCATION_BASE') or 0
+    return (self.mmio.regs.get(fb_loc_base, 0) & 0xFFFFFF) << 24

test/mockgpu/amd/amdgpu.py

@@ -127,7 +127,7 @@ class PM4Executor(AMDQueue):
     val = val_lo + (val_hi << 32)
     _ = self._next_dword() # ev
 
-    ptr = to_mv(addr_lo + (addr_hi << 32), 8)
+    ptr = to_mv(self.gpu.translate_addr(addr_lo + (addr_hi << 32)), 8)
     if mem_data_sel == 1 or mem_data_sel == 2: ptr.cast('Q')[0] = val
     elif mem_data_sel == 3:
       if mem_event_type == CACHE_FLUSH_AND_INV_TS_EVENT: ptr.cast('Q')[0] = int(time.perf_counter() * 1e8)
@@ -143,7 +143,7 @@ class PM4Executor(AMDQueue):
     dst_addr_lo = self._next_dword()
     dst_addr_hi = self._next_dword()
     assert copy_data_flags in {0x100204, 0x000204}, hex(copy_data_flags) # better fail than silently do the wrong thing
-    to_mv(dst_addr_hi<<32|dst_addr_lo, 4).cast('I')[0] = self.gpu.regs[src_addr_lo]
+    to_mv(self.gpu.translate_addr(dst_addr_hi<<32|dst_addr_lo), 4).cast('I')[0] = self.gpu.regs[src_addr_lo]
 
   def _exec_wait_reg_mem(self, n):
     assert n == 5
@@ -161,7 +161,7 @@ class PM4Executor(AMDQueue):
 
     if mem_space == 0 and mem_op == 1: mval = val # hack for memory barrier, should properly handle (req_req, reg_done)
     elif mem_space == 0: mval = self.gpu.regs[addr_hi<<32|addr_lo]
-    elif mem_space == 1: mval = to_mv(addr_lo + (addr_hi << 32), 4).cast('I')[0]
+    elif mem_space == 1: mval = to_mv(self.gpu.translate_addr(addr_lo + (addr_hi << 32)), 4).cast('I')[0]
 
     mval &= mask
 
@@ -225,7 +225,7 @@ class PM4Executor(AMDQueue):
     wptr = memoryview(bytearray(8)).cast('Q')
     rptr[0] = 0
     wptr[0] = buf_sz
-    self.ib_executor = PM4Executor(self.gpu, (addr_hi << 32) | addr_lo, buf_sz * 4, rptr, wptr)
+    self.ib_executor = PM4Executor(self.gpu, self.gpu.translate_addr((addr_hi << 32) | addr_lo), buf_sz * 4, rptr, wptr)
 
   def _exec_event_write(self, n):
     assert n == 0
@@ -276,7 +276,7 @@ class SDMAExecutor(AMDQueue):
 
   def _execute_fence(self):
     struct = sdma_pkts.fence.from_address(self.base + self.rptr[0] % self.size)
-    to_mv(struct.addr, 8).cast('Q')[0] = struct.data
+    to_mv(self.gpu.translate_addr(struct.addr), 8).cast('Q')[0] = struct.data
     self.rptr[0] += ctypes.sizeof(struct)
 
   def _execute_trap(self):
@@ -287,7 +287,7 @@ class SDMAExecutor(AMDQueue):
     struct = sdma_pkts.poll_regmem.from_address(self.base + self.rptr[0] % self.size)
 
     if struct.mem_poll == 0: mval = struct.value & struct.mask
-    elif struct.mem_poll == 1: mval = to_mv(struct.addr, 4).cast('I')[0] & struct.mask
+    elif struct.mem_poll == 1: mval = to_mv(self.gpu.translate_addr(struct.addr), 4).cast('I')[0] & struct.mask
 
     if struct.func == WAIT_REG_MEM_FUNCTION_GEQ: can_cont = bool(mval >= struct.value)
     elif struct.func == WAIT_REG_MEM_FUNCTION_EQ: can_cont = bool(mval == struct.value)
@@ -302,7 +302,7 @@ class SDMAExecutor(AMDQueue):
   def _execute_timestamp(self):
     struct = sdma_pkts.timestamp.from_address(self.base + self.rptr[0] % self.size)
 
-    mem = to_mv(struct.addr, 8).cast('Q')
+    mem = to_mv(self.gpu.translate_addr(struct.addr), 8).cast('Q')
     mem[0] = int(time.perf_counter() * 1e8)
 
     self.rptr[0] += ctypes.sizeof(struct)
@@ -313,8 +313,8 @@ class SDMAExecutor(AMDQueue):
 
   def _execute_copy(self):
     struct = sdma_pkts.copy_linear.from_address(self.base + self.rptr[0] % self.size)
-    count_cnt = to_mv(self.base + self.rptr[0] + 4, 4).cast('I')[0] & 0x3FFFFFFF
-    ctypes.memmove(struct.dst_addr, struct.src_addr, count_cnt + 1)
+    count_cnt = to_mv(self.base + self.rptr[0] % self.size + 4, 4).cast('I')[0] & 0x3FFFFFFF
+    ctypes.memmove(self.gpu.translate_addr(struct.dst_addr), self.gpu.translate_addr(struct.src_addr), count_cnt + 1)
     self.rptr[0] += ctypes.sizeof(struct)
 
 class AMDGPURegisters:
@@ -343,6 +343,7 @@ class AMDGPU(VirtGPU):
     self.queues = []
     self.arch = "cdna" if MOCKGPU_ARCH == "cdna4" else MOCKGPU_ARCH
 
+  def translate_addr(self, addr:int) -> int: return addr
   def map_range(self, vaddr, size): self.mapped_ranges.add((vaddr, size))
   def unmap_range(self, vaddr, size): self.mapped_ranges.remove((vaddr, size))
   def add_pm4_queue(self, base, size, rptr, wptr):

test/mockgpu/mockgpu.py

@@ -1,7 +1,9 @@
 import ctypes, ctypes.util, time, os, builtins, fcntl
+from tinygrad.helpers import getenv
 from tinygrad.runtime.support.hcq import FileIOInterface
 from test.mockgpu.nv.nvdriver import NVDriver
 from test.mockgpu.amd.amddriver import AMDDriver
+from test.mockgpu.am.amdriver import AMDriver, AMUSBDriver
 start = time.perf_counter()
 
 # *** ioctl lib ***
@@ -9,7 +11,8 @@ libc = ctypes.CDLL(ctypes.util.find_library("c"))
 libc.mmap.argtypes = [ctypes.c_void_p, ctypes.c_size_t, ctypes.c_int, ctypes.c_int, ctypes.c_int, ctypes.c_long]
 libc.mmap.restype = ctypes.c_void_p
 
-drivers = [AMDDriver(), NVDriver()]
+_amd_iface = getenv("AMD_IFACE", "")
+drivers = [NVDriver(), AMDriver() if _amd_iface == "PCI" else (AMUSBDriver() if _amd_iface == "USB" else AMDDriver())]
 tracked_fds = {}
 
 original_memoryview = builtins.memoryview
@@ -77,9 +80,10 @@ class MockFileIOInterface(FileIOInterface):
     return libc.mmap(start, sz, prot, flags, self.fd, offset)
 
   def read(self, size=None, binary=False, offset=None):
-    if binary: raise NotImplementedError()
     if self.fd in tracked_fds:
+      if offset is not None: tracked_fds[self.fd].seek(offset)
       return tracked_fds[self.fd].read_contents(size)
+    if binary: raise NotImplementedError()
     with open(self.fd, "rb" if binary else "r", closefd=False) as file:
       if file.tell() >= os.fstat(self.fd).st_size: file.seek(0)
       return file.read(size)
@@ -89,13 +93,20 @@ class MockFileIOInterface(FileIOInterface):
       return tracked_fds[self.fd].list_contents()
     return os.listdir(self.path)
 
-  def write(self, content, binary=False, offset=None): raise NotImplementedError()
+  def write(self, content, binary=False, offset=None):
+    if self.fd in tracked_fds:
+      if offset is not None: tracked_fds[self.fd].seek(offset)
+      return tracked_fds[self.fd].write_contents(content)
+    raise NotImplementedError()
   def seek(self, offset):
     if self.fd in tracked_fds:
       tracked_fds[self.fd].seek(offset)
     else:
       os.lseek(self.fd, offset, os.SEEK_CUR)
   @staticmethod
+  def anon_mmap(start, sz, prot, flags, offset):
+    return FileIOInterface._mmap(start, sz, prot, flags & ~0x4a000, -1, offset)  # strip MAP_LOCKED|MAP_POPULATE|MAP_HUGETLB
+  @staticmethod
   def exists(path): return _open(path, os.O_RDONLY) is not None
   @staticmethod
   def readlink(path): raise NotImplementedError()

test/mockgpu/usb.py

@@ -1,16 +1,213 @@
+from __future__ import annotations
+import ctypes, mmap, struct, sys
+if sys.platform != "win32": from tinygrad.runtime.autogen import libc
+
 class MockUSB:
   def __init__(self, mem):
     self.mem = mem
-
-  def read(self, address, size):
-    return bytes(self.mem[address:address+size])
-
-  def write(self, address, data, ignore_cache=False):
-    self.mem[address:address+len(data)] = data
-
+  def read(self, address, size): return bytes(self.mem[address:address+size])
+  def write(self, address, data, ignore_cache=False): self.mem[address:address+len(data)] = data
   def pcie_mem_req(self, address, value=None, size=1):
     if value is None: return int.from_bytes(self.mem[address:address+size], "little")
     else: self.mem[address:address+size] = value.to_bytes(size, "little")
-
   def pcie_mem_write(self, address, values, size):
     for i, value in enumerate(values): self.pcie_mem_req(address + i * size, value, size)
+
+# *** ASM24 Controller Mock ***
+
+_mock_usb_state: MockASM24State|None = None
+
+class MockASM24State:
+  """Mock ASM24 controller: XRAM memory map, DMA windows, TLP engine, PCI config space.
+
+  Memory map (64KB XRAM):
+    0xA000-0xAFFF: DMA window -> sys 0x820000
+    0xB000-0xB1FF: DMA window -> sys 0x800000
+    0xB200-0xB7FF: PCI MMIO (TLP engine)
+    0xF000-0xFFFF: DMA window -> sys 0x200000 (512KB)
+  """
+  XRAM_SIZE = 0x10000
+
+  TLP_FMT_TYPE = 0xB210
+  TLP_BYTE_EN = 0xB217
+  TLP_ADDR_LO = 0xB218
+  TLP_ADDR_HI = 0xB21C
+  TLP_DATA = 0xB220
+  TLP_COMPL = 0xB22A
+  TLP_TRIGGER = 0xB254
+  TLP_LINK_STATUS = 0xB284
+  TLP_STATUS = 0xB296
+
+  def __init__(self, gpu, driver, vram_size:int, doorbell_size:int, mmio_size:int):
+    self.gpu, self.driver = gpu, driver
+    self._xram = bytearray(self.XRAM_SIZE)
+
+    self._doorbell_addr = libc.mmap(0, doorbell_size, mmap.PROT_READ | mmap.PROT_WRITE, mmap.MAP_SHARED, gpu.doorbell_fd, 0)
+    self._doorbell = (ctypes.c_ubyte * doorbell_size).from_address(self._doorbell_addr)
+
+    # DMA windows: ctrl_addr -> (host_addr, size)
+    self._dma_regions: dict[int, tuple[int, int]] = {}
+    self._add_dma_window(0xF000, 0x200000, 0x80000)
+    self._add_dma_window(0xA000, 0x820000, 0x1000)
+    self._add_dma_window(0xB000, 0x800000, 0x200)
+
+    # PCI config space: (bus,dev,fn) -> bytearray(4096)
+    self._pci_cfg: dict[tuple[int,int,int], bytearray] = {}
+
+    # GPU BAR definitions: reg_offset -> (size, type_bits, is_64bit)
+    self._gpu_bars: dict[int, tuple[int, int, bool]] = {
+      0x10: (vram_size, 0x0C, True),       # BAR0: VRAM, 64-bit prefetchable
+      0x18: (doorbell_size, 0x00, False),   # BAR2: doorbell, 32-bit
+      0x1C: (0, 0x00, False),              # BAR3: unused
+      0x20: (0, 0x00, False),              # BAR4: unused
+      0x24: (mmio_size, 0x00, False),      # BAR5: MMIO, 32-bit
+    }
+    self._bar_addrs: dict[int, tuple[int, int]] = {}  # reg_offset -> (addr, size)
+
+    # Initialize GPU config space (bus=4, dev=0, fn=0) with BAR type bits and REBAR capability
+    gpu_cfg = self._get_cfg(4, 0, 0)
+    for reg_off, (sz, type_bits, _) in self._gpu_bars.items():
+      if sz > 0: struct.pack_into('<I', gpu_cfg, reg_off, type_bits)
+    struct.pack_into('<I', gpu_cfg, 0x100, 0x15 | (1 << 16))  # REBAR cap header: id=0x15, version=1, next=0
+    struct.pack_into('<I', gpu_cfg, 0x104, sum(1 << (i + 4) for i in range(10)))  # supported sizes up to 512MB
+
+  def _get_cfg(self, bus:int, dev:int, fn:int) -> bytearray:
+    if (key:=(bus, dev, fn)) not in self._pci_cfg: self._pci_cfg[key] = bytearray(4096)
+    return self._pci_cfg[key]
+
+  def _add_dma_window(self, ctrl_addr:int, sys_addr:int, size:int):
+    host_addr = libc.mmap(0, size, mmap.PROT_READ | mmap.PROT_WRITE, mmap.MAP_SHARED | mmap.MAP_ANONYMOUS, -1, 0)
+    self._dma_regions[ctrl_addr] = (host_addr, size)
+    for off in range(0, size, 0x1000): self.gpu._sysmem_map[sys_addr + off] = host_addr + off
+
+  # --- XRAM access ---
+
+  def _xram_read(self, addr:int, length:int) -> bytes:
+    for ctrl_addr, (host_addr, dma_size) in self._dma_regions.items():
+      if ctrl_addr <= addr < ctrl_addr + dma_size:
+        return bytes((ctypes.c_ubyte * length).from_address(host_addr + (addr - ctrl_addr)))
+    return bytes(self._xram[addr:addr+length])
+
+  def _xram_write_byte(self, addr:int, value:int):
+    for ctrl_addr, (host_addr, dma_size) in self._dma_regions.items():
+      if ctrl_addr <= addr < ctrl_addr + dma_size:
+        (ctypes.c_ubyte * 1).from_address(host_addr + (addr - ctrl_addr))[0] = value
+        return
+    if addr == self.TLP_STATUS:
+      self._xram[addr] &= ~value & 0xFF
+      return
+    self._xram[addr] = value
+    if addr == self.TLP_TRIGGER and value == 0x0F: self._process_tlp()
+
+  # --- TLP engine ---
+
+  def _process_tlp(self):
+    fmt_type, byte_en = self._xram[self.TLP_FMT_TYPE], self._xram[self.TLP_BYTE_EN]
+    addr_lo = int.from_bytes(self._xram[self.TLP_ADDR_LO:self.TLP_ADDR_LO+4], 'big')
+    addr_hi = int.from_bytes(self._xram[self.TLP_ADDR_HI:self.TLP_ADDR_HI+4], 'big')
+    address = addr_lo | (addr_hi << 32)
+
+    size, offset, tmp = 0, 0, byte_en
+    while tmp and not (tmp & 1):
+      offset += 1
+      tmp >>= 1
+    while tmp:
+      size += tmp & 1
+      tmp >>= 1
+
+    is_write, is_cfg = bool(fmt_type & 0x40), (fmt_type & 0xbe) == 0x04
+
+    if is_cfg:
+      bus, dev, fn, byte_addr = (address >> 24) & 0xFF, (address >> 19) & 0x1F, (address >> 16) & 0x7, address & 0xFFC
+      if is_write:
+        data = int.from_bytes(self._xram[self.TLP_DATA:self.TLP_DATA+4], 'big')
+        self._cfg_write(bus, dev, fn, byte_addr + offset, (data >> (8 * offset)) & ((1 << (8 * size)) - 1), size)
+      else:
+        self._xram[self.TLP_DATA:self.TLP_DATA+4] = int.from_bytes(self._get_cfg(bus, dev, fn)[byte_addr:byte_addr+4], 'little').to_bytes(4, 'big')
+      self._xram[self.TLP_COMPL:self.TLP_COMPL+2] = (4).to_bytes(2, 'big')
+      self._xram[self.TLP_LINK_STATUS] = 0x01 if not is_write else 0x00
+      self._xram[self.TLP_STATUS] = 0x02
+      return
+
+    if is_write:
+      data = int.from_bytes(self._xram[self.TLP_DATA:self.TLP_DATA+4], 'big')
+      self._pcie_dispatch(address + offset, (data >> (8 * offset)) & ((1 << (8 * size)) - 1), size)
+    else:
+      result = self._pcie_dispatch(address + offset, None, size)
+      if result is not None:
+        self._xram[self.TLP_DATA:self.TLP_DATA+4] = ((result << (8 * offset)) & 0xFFFFFFFF).to_bytes(4, 'big')
+
+    self._xram[self.TLP_COMPL:self.TLP_COMPL+2] = (size & 0xFFF).to_bytes(2, 'big')
+    self._xram[self.TLP_LINK_STATUS] = 0x01 if not is_write else 0x00
+    self._xram[self.TLP_STATUS] = 0x02
+
+  def _cfg_write(self, bus:int, dev:int, fn:int, byte_addr:int, val:int, size:int):
+    cfg = self._get_cfg(bus, dev, fn)
+
+    # Handle BAR register writes for GPU device (bus=4, dev=0, fn=0)
+    if (bus, dev, fn) == (4, 0, 0) and 0x10 <= byte_addr < 0x28 and size == 4:
+      reg_off = byte_addr & ~0x3
+      if (bar_def:=self._gpu_bars.get(reg_off)) is not None:
+        bar_size, type_bits, is_64 = bar_def
+        if bar_size == 0: return  # unused BAR
+        if val == 0xFFFFFFFF:  # size probe
+          struct.pack_into('<I', cfg, reg_off, (~(bar_size - 1)) & 0xFFFFFFF0 | type_bits)
+        else:
+          struct.pack_into('<I', cfg, reg_off, val)
+          hi = struct.unpack_from('<I', cfg, reg_off + 4)[0] if is_64 else 0
+          self._bar_addrs[reg_off] = ((hi << 32) | (val & ~0xF), bar_size)
+        return
+      # Check if upper 32 bits of a 64-bit BAR
+      for breg, (bsz, _, b64) in self._gpu_bars.items():
+        if b64 and reg_off == breg + 4:
+          struct.pack_into('<I', cfg, reg_off, 0xFFFFFFFF if val == 0xFFFFFFFF else val)
+          if val != 0xFFFFFFFF:
+            self._bar_addrs[breg] = ((val << 32) | (struct.unpack_from('<I', cfg, breg)[0] & ~0xF), bsz)
+          return
+
+    # Generic config write
+    for i in range(size): cfg[byte_addr + i] = (val >> (8 * i)) & 0xFF
+
+  def _pcie_dispatch(self, address:int, value:int|None, size:int) -> int|None:
+    for reg_off, (bar_addr, bar_size) in self._bar_addrs.items():
+      if bar_addr <= address < bar_addr + bar_size:
+        offset = address - bar_addr
+        if reg_off == 0x10:  # BAR0 - VRAM
+          if value is None: return int.from_bytes(bytes(self.gpu.vram[offset:offset+size]), "little")
+          self.gpu.vram[offset:offset+size] = list(value.to_bytes(size, "little"))
+          return None
+        if reg_off == 0x18:  # BAR2 - Doorbell
+          if value is None: return int.from_bytes(bytes(self._doorbell[offset:offset+size]), "little")
+          for i, b in enumerate(value.to_bytes(size, "little")): self._doorbell[offset + i] = b
+          self.driver._emulate_execute()
+          return None
+        if reg_off == 0x24:  # BAR5 - MMIO
+          if value is None: return self.gpu.mmio[offset // 4]
+          self.gpu.mmio[offset // 4] = value
+          return None
+    raise ValueError(f"PCIe address {address:#x} not mapped to any BAR")
+
+  # --- CDB processing (called by MockUSB3.send_batch) ---
+
+  def process_cdb(self, cdb:bytes, rlen:int, send_data:bytes|None) -> bytes|None:
+    op = cdb[0]
+    if op == 0xE5:  # write byte
+      self._xram_write_byte(((cdb[2] << 16) | (cdb[3] << 8) | cdb[4]) & 0xFFFF, cdb[1])
+      return None
+    if op == 0xE4:  # read
+      return self._xram_read(((cdb[2] << 16) | (cdb[3] << 8) | cdb[4]) & 0xFFFF, cdb[1])
+    if op == 0x8A and send_data is not None and 0xF000 in self._dma_regions:  # SCSI write
+      host_addr, dma_size = self._dma_regions[0xF000]
+      ctypes.memmove(host_addr, send_data, min(len(send_data), dma_size))
+    return None
+
+class MockUSB3:
+  def __init__(self, *args, **kwargs): pass
+  def send_batch(self, cdbs:list[bytes], idata:list[int]|None=None, odata:list[bytes|None]|None=None) -> list[bytes|None]:
+    assert _mock_usb_state is not None
+    idata, odata = idata or [0] * len(cdbs), odata or [None] * len(cdbs)
+    results: list[bytes|None] = []
+    for cdb, rlen, sdata in zip(cdbs, idata, odata):
+      result = _mock_usb_state.process_cdb(cdb, rlen, sdata)
+      results.append(result if rlen > 0 else None)
+    return results

test/null/test_graph_rewrite.py

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

test/null/test_multitensor.py

@@ -38,6 +38,27 @@ class TestMultiRamUsage(unittest.TestCase):
     _ = Tensor.zeros(self.N, self.N).contiguous().shard(devices_2, axis=0).contiguous().realize()
     self.assertUsed(self.N*self.N*4) # sharding should not increase total ram usage
 
+  def test_sharded_memory_replicated(self):
+    devices_4 = tuple(f"NULL:{i+1}" for i in range(4))
+    X = Tensor.ones(256).contiguous().realize()
+    self.assertUsed(256 * 4)
+    X.shard_(devices_4).realize()
+    self.assertUsed(256 * 4 * 4)
+
+  def test_sharded_memory_replicated_const(self):
+    devices_4 = tuple(f"NULL:{i+1}" for i in range(4))
+    X = Tensor.ones(256).realize()
+    self.assertUsed(0)
+    X.shard_(devices_4).realize()
+    self.assertUsed(256 * 4 * 4)  # TODO: can be zero
+
+  def test_sharded_memory_axis_const(self):
+    devices_4 = tuple(f"NULL:{i+1}" for i in range(4))
+    X = Tensor.ones(256).realize()
+    self.assertUsed(0)
+    X.shard_(devices_4, axis=0).realize()
+    self.assertUsed(256 * 4)  # TODO: can be zero
+
   def _test_matmul_half(self, dev_count:int):
     N = 32
     total_mem = {}
@@ -66,5 +87,13 @@ class TestMultiAxis(unittest.TestCase):
     self.assertEqual(t.reshape(2, 16).uop.axis, 0)
     self.assertEqual(t.reshape(2, 2, 8).uop.axis, 0)
 
+  def test_empty_like_sharded(self):
+    t = Tensor.ones(4, 8).shard(("NULL:0", "NULL:1"), axis=0)
+    e = t.empty_like()
+    self.assertEqual(e.shape, t.shape)
+    self.assertEqual(e.device, t.device)
+    self.assertEqual(e.uop.axis, 0)
+    self.assertTrue(e.uop.has_buffer_identity())
+
 if __name__ == '__main__':
   unittest.main()

test/null/test_real_world.py

@@ -98,7 +98,7 @@ class TestRealWorld(unittest.TestCase):
     @TinyJit
     def test(t, v):
       with Context(JIT=0): return model(t, v).realize()
-    helper_test("test_gpt2", lambda: (Tensor([[1,]]),Variable("pos", 1, 100).bind(1)), test, 0.23, 160, all_jitted=True)
+    helper_test("test_gpt2", lambda: (Tensor([[1,]]),Variable("pos", 1, 100).bind(1)), test, 0.23, 168, all_jitted=True)
 
   @slow
   def test_train_mnist(self):

test/null/test_rewrite_map.py

@@ -1,202 +0,0 @@
-import unittest
-from tinygrad import dtypes
-from tinygrad.uop.ops import UOp, graph_rewrite_map, _substitute
-from tinygrad.uop.symbolic import symbolic
-
-class TestRewriteMap(unittest.TestCase):
-  def test_substitute(self):
-    a = UOp.variable('a', 0, 10)
-    b = UOp.variable('b', 0, 10)
-    c = UOp.variable('c', 0, 10)
-    e = UOp.variable('e', 0, 10)
-    ret = (a+b)*c
-    sub = {a+b: e}
-    sub_map = graph_rewrite_map(ret, _substitute, sub, bottom_up=True)
-    self.assertIs(sub_map[a+b], e)
-    self.assertIs(sub_map[(a+b)*c], e*c)
-
-  def test_substitute_depth_2(self):
-    a = UOp.variable('a', 0, 10)
-    b = UOp.variable('b', 0, 10)
-    c = UOp.variable('c', 0, 10)
-    d = UOp.variable('d', 0, 10)
-    e = UOp.variable('e', 0, 10)
-    f = UOp.variable('f', 0, 10)
-    ret = (a+b)*c+d
-    sub = {a+b: e, (a+b)*c: f}
-    sub_map = graph_rewrite_map(ret, _substitute, sub, bottom_up=True)
-    self.assertIs(sub_map[a+b], e)
-    self.assertIs(sub_map[(a+b)*c], f)
-
-  def test_multistage_substitute(self):
-    a = UOp.variable('a', 0, 10)
-    b = UOp.variable('b', 0, 10)
-    c = UOp.variable('c', 0, 10)
-    d = UOp.variable('d', 0, 10)
-    sub1 = {a+b:c}
-    start = (a+b)*c
-    # stage 1: (a+b)*c -> c*c
-    sub_map1 = graph_rewrite_map(start, _substitute, sub1, bottom_up=True)
-    self.assertIs(sub_map1[(a+b)*c], c*c)
-    # stage 2: c*c -> d
-    sub2 = {c*c:d}
-    sub_map2 = graph_rewrite_map(sub_map1[start], _substitute, sub2, input_map=sub_map1, bottom_up=True)
-    # (a+b)*c -> c*c -> d
-    self.assertIs(sub_map2[(a+b)*c], d)
-
-  def test_add_zero(self):
-    # Build a small graph: add(0, add(const=0, const=5))
-    zero_node = UOp.const(dtypes.index, 0)
-    five_node = UOp.const(dtypes.index, 5)
-    inner_add = zero_node + five_node
-    root_add = zero_node + inner_add
-
-    # Perform top-down rewrite
-    node_map = graph_rewrite_map(root_add, symbolic)
-
-    # We expect that add(0, add(0, 5)) -> add(0, 5) -> 5
-    # Check the mapping
-    assert node_map[root_add] == five_node
-    assert node_map[inner_add] == five_node
-    # zero_node and five_node map to themselves
-    assert node_map[zero_node] == zero_node
-    assert node_map[five_node] == five_node
-
-  def test_double_neg(self):
-    """
-    Test rewriting neg(neg(5)) => 5 using symbolic.
-    """
-    # In some versions of TinyGrad, you might do: (-(-five_node))
-    five_node = UOp.const(dtypes.index, 5)
-    # If your code allows UOp(...), do that; else you might do something like:
-    # double_neg_five = -(-five_node)
-    # But let's be explicit:
-    neg_five = -five_node
-    double_neg_five = -neg_five
-
-    node_map = graph_rewrite_map(double_neg_five, symbolic)
-
-    # node_map should map double_neg_five -> five_node
-    self.assertEqual(node_map[double_neg_five], five_node)
-    # five_node maps to itself
-    self.assertEqual(node_map[five_node], five_node)
-
-  def test_add_zero_and_double_neg(self):
-    """
-    Combine both rewrites: add(0, neg(neg(5))) => add(0, 5) => 5
-    """
-    zero_node = UOp.const(dtypes.index, 0)
-    five_node = UOp.const(dtypes.index, 5)
-    neg_five = -five_node
-    double_neg_five = -neg_five
-    root_add = zero_node + double_neg_five
-
-    node_map = graph_rewrite_map(root_add, symbolic)
-
-    # node_map: root_add -> five_node, double_neg_five -> five_node
-    self.assertEqual(node_map[root_add], five_node)
-    self.assertEqual(node_map[double_neg_five], five_node)
-    # zero_node, five_node map to themselves
-    self.assertEqual(node_map[zero_node], zero_node)
-    self.assertEqual(node_map[five_node], five_node)
-
-  def test_multi_var_rewrites(self):
-    x_var = UOp.variable('x', 0, 10)
-    y_var = UOp.variable('y', -5, 5)
-    zero_node = UOp.const(dtypes.index, 0)
-
-    sum_with_zero = y_var + zero_node    # (y + 0)
-    combined = x_var + sum_with_zero     # x + (y + 0)
-    double_neg = -(-combined)           # neg(neg(x + y))
-    final_expr = zero_node + double_neg  # 0 + (x + y)
-
-    node_map = graph_rewrite_map(final_expr, symbolic)
-
-    # The final root should be (x_var + y_var).
-    expected = x_var + y_var
-
-    # Each sub-expression has its own "final" result.
-    # (y + 0) -> y_var
-    self.assertEqual(node_map[sum_with_zero], y_var)
-    # (x + (y+0)) -> (x + y)
-    self.assertEqual(node_map[combined], expected)
-    # neg(neg(x+y)) -> (x + y)
-    self.assertEqual(node_map[double_neg], expected)
-    # 0 + (x+y) -> (x + y)
-    self.assertEqual(node_map[final_expr], expected)
-
-    # x_var, y_var, zero_node remain unchanged
-    self.assertEqual(node_map[x_var], x_var)
-    self.assertEqual(node_map[y_var], y_var)
-    self.assertEqual(node_map[zero_node], zero_node)
-
-  def test_complex_multi_var_edges(self):
-    """
-    Build a multi-variable expression with multiple intermediates:
-
-      x_var = UOp.variable('x', 1, 10)
-      y_var = UOp.variable('y', -5, 5)
-      z_var = UOp.variable('z', 0, 5)
-      zero_node = UOp.const(dtypes.int, 0)
-      one_node = UOp.const(dtypes.int, 1)
-
-      yz_sum       = y_var + z_var
-      yz_sum_zero  = yz_sum + zero_node   -> rewrites to yz_sum
-      yz_neg       = -yz_sum_zero        -> -(y+z)
-      yz_dneg      = -yz_neg             -> y+z    (double neg gone)
-      x_plus_yz    = x_var + yz_dneg     -> x + (y+z)
-      double_neg_x = -(-x_plus_yz)       -> x + (y+z)
-      final_expr   = double_neg_x * one_node -> x + (y+z)
-
-    We expect the final result to be (x + (y+z)).
-    Each original node should map to the final node that replaces it,
-    which might be structurally equivalent but not the same reference.
-    """
-    x_var = UOp.variable('x', 1, 10)
-    y_var = UOp.variable('y', -5, 5)
-    z_var = UOp.variable('z', 0, 5)
-    zero_node = UOp.const(dtypes.index, 0)
-    one_node = UOp.const(dtypes.index, 1)
-
-    # Build sub-expressions
-    yz_sum = y_var + z_var              # (y + z)
-    yz_sum_zero = yz_sum + zero_node    # (y + z) + 0
-    yz_neg = -yz_sum_zero               # -(y+z)
-    yz_dneg = -yz_neg                   # -(-(y+z)) -> (y+z)
-    x_plus_yz = x_var + yz_dneg         # x + (y+z)
-    double_neg_x = -(-x_plus_yz)        # neg(neg(x+(y+z))) -> x+(y+z)
-    final_expr = double_neg_x * one_node  # (x+(y+z)) * 1 -> x+(y+z)
-
-    node_map = graph_rewrite_map(final_expr, symbolic)
-
-    # (y + z) is unchanged
-    self.assertEqual(node_map[yz_sum], yz_sum)
-
-    # (y+z) + 0 => (y+z)
-    self.assertEqual(node_map[yz_sum_zero], yz_sum)
-
-    # -(y+z) remains -(y+z), but might be a new UOp with updated children
-    # Compare structurally to -(y_var + z_var).
-    self.assertEqual(node_map[yz_neg], -yz_sum)
-
-    # -(-(y+z)) => (y+z)
-    self.assertEqual(node_map[yz_dneg], yz_sum)
-
-    # x + (y+z) => might get recreated if yz_dneg was changed, so compare to x + yz_sum
-    self.assertEqual(node_map[x_plus_yz], x_var + yz_sum)
-
-    # -(-(x+(y+z))) => x + (y+z)
-    self.assertEqual(node_map[double_neg_x], x_var + yz_sum)
-
-    # (x+(y+z)) * 1 => x+(y+z)
-    self.assertEqual(node_map[final_expr], x_var + yz_sum)
-
-    # Unchanged atomic nodes map to themselves
-    self.assertEqual(node_map[x_var], x_var)
-    self.assertEqual(node_map[y_var], y_var)
-    self.assertEqual(node_map[z_var], z_var)
-    self.assertEqual(node_map[zero_node], zero_node)
-    self.assertEqual(node_map[one_node], one_node)
-
-if __name__ == "__main__":
-  unittest.main()

test/null/test_schedule.py

@@ -169,7 +169,7 @@ class TestSchedule(unittest.TestCase):
   def test_empty_is_not_realized(self):
     a = Tensor.empty(10)
     child = a+2
-    assert a.uop.is_realized
+    assert not a.uop.is_realized
     child.realize()
     assert a.uop.is_realized
 
@@ -185,7 +185,7 @@ class TestSchedule(unittest.TestCase):
   def test_childless_empty_never_allocates(self):
     a = Tensor.empty(10)
     a.realize()
-    assert not a.uop.buffer.is_allocated()
+    assert not a.uop.is_realized
 
   def test_simplify_padded_const(self):
     a, _ = Tensor.empty(1022).cummax(axis=0)
@@ -412,20 +412,20 @@ class TestSchedule(unittest.TestCase):
       out = bn(c1(img)).relu()
       check_schedule(out, 4, [c1.weight, c1.bias])
 
-  def test_fold_conv_batchnorm_optim(self):
-    # this is too high
-    for optim, cnt in [(nn.optim.Adam, 27), (nn.optim.SGD, 7)]:
-      with self.subTest(optim=optim.__name__):
-        with Tensor.train():
-          img = Tensor.ones(1,3,4,4)
-          c1 = nn.Conv2d(3,32,3)
-          bn = nn.BatchNorm2d(32, track_running_stats=False)
-          _realize_weights([c1, bn])
-          opt = optim(nn.state.get_parameters([c1, bn]))
-          img_bn = bn(c1(img)).elu().sum()
-          opt.zero_grad()
-          img_bn.backward()
-          check_schedule(opt.schedule_step(), cnt)
+  def test_fold_conv_batchnorm_optim(self, adam=False):
+    # 2 is too low?
+    optim, cnt = (nn.optim.Adam, 16) if adam else (nn.optim.SGD, 2)
+    with Tensor.train():
+      img = Tensor.ones(1,3,4,4)
+      c1 = nn.Conv2d(3,32,3)
+      bn = nn.BatchNorm2d(32, track_running_stats=False)
+      _realize_weights([c1, bn])
+      opt = optim(nn.state.get_parameters([c1, bn]))
+      img_bn = bn(c1(img)).elu().sum()
+      opt.zero_grad()
+      img_bn.backward()
+      check_schedule(opt.schedule_step(), cnt)
+  def test_fold_conv_batchnorm_optim_adam(self): self.test_fold_conv_batchnorm_optim(True)
 
   def test_fold_batchnorm_backward(self):
     with Tensor.train():
@@ -774,7 +774,7 @@ class TestSchedule(unittest.TestCase):
       _realize_weights(layer)
       opt = nn.optim.Adam(nn.state.get_parameters(layer), lr=1e-4)
       layer(x).relu().sum().backward()
-      check_schedule(opt.schedule_step(), 19)
+      check_schedule(opt.schedule_step(), 13)
 
   def test_adam_conv_fuse(self):
     with Tensor.train():
@@ -784,7 +784,7 @@ class TestSchedule(unittest.TestCase):
       opt = nn.optim.Adam(nn.state.get_parameters(c1), lr=1e-4)
       opt.zero_grad()
       c1(img).relu().sum().backward()
-      check_schedule(opt.schedule_step(), 19)
+      check_schedule(opt.schedule_step(), 13)
 
   def test_adam_2convs_fuse(self):
     with Tensor.train():
@@ -795,7 +795,7 @@ class TestSchedule(unittest.TestCase):
       opt = nn.optim.Adam(nn.state.get_parameters([c1, c2]), lr=1e-4)
       opt.zero_grad()
       c2(c1(img).relu()).relu().sum().backward()
-      check_schedule(opt.schedule_step(), 21)
+      check_schedule(opt.schedule_step(), 15)
 
   def test_sgd_conv_fuse(self):
     with Tensor.train():
@@ -827,7 +827,7 @@ class TestSchedule(unittest.TestCase):
       opt = nn.optim.SGD(nn.state.get_parameters([c1, c2]), nesterov=True, momentum=0.9, weight_decay=0.1)
       opt.zero_grad()
       c2(c1(img).relu()).relu().sum().backward()
-      check_schedule(opt.schedule_step(), 13)
+      check_schedule(opt.schedule_step(), 11)
 
   def test_sgd_4convs_fuse(self):
     with Tensor.train():

test/null/test_tensor_metadata.py

@@ -4,6 +4,7 @@ from tinygrad.tensor import _METADATA
 from tinygrad.engine.realize import capturing
 from tinygrad.helpers import Context
 
+@unittest.skip("tensor metadata is no longer supported")
 class TestTensorMetadata(unittest.TestCase):
   def setUp(self) -> None:
     _METADATA.set(None)

test/null/test_uop_symbolic.py

@@ -390,6 +390,9 @@ class TestSymbolic(unittest.TestCase):
     self.helper_test_variable(Variable("a", 0, 6) < 3, 0, 1, "(a<3)")
     self.helper_test_variable(Variable("a", 0, 6) < 8, 1, 1, "True")
 
+  def test_cast_bool(self):
+    self.helper_test_variable(Variable("a", 0, 10).cast(dtypes.bool), 0, 1, "a!=0")
+
   def test_lt_sum_remove(self):
     self.helper_test_variable(Variable("a", 0, 6) + 2 < 3, 0, 1, "(a<1)")
 
@@ -655,6 +658,24 @@ class TestSymbolic(unittest.TestCase):
     with self.assertRaises(AssertionError):
       self.helper_test_variable((31 * b + 1) % 18 + ((31 * b + 1) // 18) * 18, 1, 3101, "((b*31)+1)")
 
+  def test_div_mod_recombine_3level(self):
+    gidx = Variable("gidx", 0, 150527)
+    self.helper_test_variable(gidx//3%224*3 + gidx%3 + gidx//672*672, 0, 150527, "gidx")
+    # different shapes
+    x = Variable("x", 0, 5*7*11-1)
+    self.helper_test_variable(x//11%7*11 + x%11 + x//77*77, 0, 5*7*11-1, "x")
+    # result is x//a*c2 not just x
+    x2 = Variable("x2", 0, 5*6*7-1)
+    self.helper_test_variable(x2//7%6*14 + x2//42*84, 0, (5*6*7-1)//7*14, "(x2//7*14)")
+    # negative variable range
+    xn = Variable("x", -1000, 1000)
+    self.helper_test_variable(xn//3%224*3 + xn%3 + xn//672*672, -1000, 1000, "x")
+    self.helper_test_variable(xn//3%7*3 + xn//21*21, -999, 999, "(x//3*3)")
+    # should NOT simplify: a*c1 != b (3*224 != 600)
+    self.helper_test_variable(gidx//3%224*3 + gidx//600*600, 0, 150669, "(gidx//600*600+gidx//3%224*3)")
+    # should NOT simplify: c1*c2 != c3 (224*3 != 700)
+    self.helper_test_variable(gidx//3%224*3 + gidx//672*700, 0, 156769, "(gidx//672*700+gidx//3%224*3)")
+
   def test_div_mod_recombine_with_gcd(self):
     b = Variable("b", 0, 100)
     exp = (16 * b + 2) % 18 + ((16 * b + 2) // 18) * 18

test/null/test_viz.py

@@ -282,9 +282,10 @@ class TestVizIntegration(BaseTestViz):
     ast = Tensor.schedule(Tensor.empty(4)+Tensor.empty(4))[0].ast
     prg = get_program(ast, Device[Device.DEFAULT].renderer)
     lst = get_viz_list()
-    self.assertEqual(len(lst), 2)
-    self.assertEqual(lst[0]["name"], "Schedule 1 Kernel n1")
-    self.assertEqual(lst[1]["name"], prg.name)
+    self.assertEqual(len(lst), 3)
+    self.assertEqual(lst[0]["name"], "Process 1 Buffer n1")
+    self.assertEqual(lst[1]["name"], "Schedule 1 Kernel n1")
+    self.assertEqual(lst[2]["name"], prg.name)
 
   # schedule graph CALL nodes have a link to jump to codegen
   def test_link_sched_codegen(self):
@@ -293,8 +294,9 @@ class TestVizIntegration(BaseTestViz):
     sched = Tensor.schedule(c1, c2)
     prgs = [si.lower().prg.p.name for si in sched]
     lst = get_viz_list()
-    viz_kernel = next(i for i,s in enumerate(lst[0]["steps"]) if s["name"] == "View Kernel Graph")
-    graph = next(get_viz_details(0, viz_kernel))["graph"]
+    sched_idx = next(i for i,l in enumerate(lst) if l["name"].startswith("Schedule"))
+    viz_kernel = next(i for i,s in enumerate(lst[sched_idx]["steps"]) if s["name"] == "View Kernel Graph")
+    graph = next(get_viz_details(sched_idx, viz_kernel))["graph"]
     call_nodes = [n for n in graph.values() if n["label"].startswith("CALL")]
     for i,n in enumerate(call_nodes):
       assert n["ref"] is not None

test/null/test_winograd.py

@@ -45,7 +45,7 @@ class TestWinograd(unittest.TestCase):
 
     # TODO: what's optimal on this?
     self.assertLess(ops_ratio, 4.3)
-    self.assertLess(mem_ratio, 3)
+    self.assertLess(mem_ratio, 4)
 
   def test_dtype(self):
     IC, OC, X, Y = 4,4,9,9

test/testextra/test_cfg_viz.py

@@ -69,6 +69,8 @@ class TestCfg(unittest.TestCase):
     self.assertEqual(len(references["r0"]), 2)
     insts = [cfg["pc_tokens"][pc][0]["st"] for pc in references["r0"]]
     self.assertEqual(insts, ['s_mov_b32', 's_cmp_eq_u64'])
+    end_block_content = "\n".join(" ".join(t["st"] for t in cfg["pc_tokens"][pc]) for pc in list(cfg["blocks"].values())[-1])
+    self.assertEqual(end_block_content, "s_endpgm\ns_code_end (217x)")
 
   def test_loop(self):
     k = Kernel(arch=Device["AMD"].arch)

test/testextra/test_hk_fa.py

@@ -128,7 +128,7 @@ class TestFA(unittest.TestCase):
     assert_allclose(k.grad, k_ref.grad, atol=1e-5, rtol=1e-5)
     assert_allclose(v.grad, v_ref.grad, atol=1e-5, rtol=1e-5)
 
-  def test_fast_fa_bwd_multidevice(self):
+  def test_fast_fa_bwd_dp(self):
     Tensor.manual_seed(42)
 
     B, N, H, H_KV, D = 2, 1024, 32, 8, 128
@@ -175,5 +175,52 @@ class TestFA(unittest.TestCase):
     assert_allclose(v.grad, v_ref.grad, atol=1e-5, rtol=1e-5)
     assert_allclose(k.grad, k_ref.grad, atol=1e-5, rtol=1e-5)
 
+  def test_fast_fa_bwd_mp(self):
+    Tensor.manual_seed(42)
+
+    B, N, H, H_KV, D = 2, 1024, 32, 8, 128
+    GPUS = tuple(f"AMD:{i}" for i in range(B))
+
+    with Context(DEBUG=0):
+      base_q = Tensor.randn(B, N, H, D, dtype=dtypes.bfloat16, requires_grad=True).contiguous()
+      base_k = Tensor.randn(B, N, H_KV, D, dtype=dtypes.bfloat16, requires_grad=True).contiguous()
+      base_v = Tensor.randn(B, N, H_KV, D, dtype=dtypes.bfloat16, requires_grad=True).contiguous()
+
+      base_do = Tensor.ones(B, N, H, D, dtype=dtypes.float32).contiguous()
+
+    with Context(DEBUG=0):
+      q = base_q.clone().requires_grad_(True).shard(GPUS, axis=2)
+      k = base_k.clone().requires_grad_(True).shard(GPUS, axis=2)
+      v = base_v.clone().requires_grad_(True).shard(GPUS, axis=2)
+      Tensor.realize(q, k, v)
+
+      do = base_do.clone().shard(GPUS, axis=2)
+      Tensor.realize(do)
+
+    q_, k_, v_ = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)
+    out = flash_attention(q_, k_, v_, is_causal=True)
+    out = out.float().transpose(1, 2)
+    out.backward(do)
+    Tensor.realize(q.grad, k.grad, v.grad)
+
+    with Context(DEBUG=0):
+      q_ref = base_q.clone().requires_grad_(True)
+      k_ref = base_k.clone().requires_grad_(True)
+      v_ref = base_v.clone().requires_grad_(True)
+      Tensor.realize(q_ref, k_ref, v_ref)
+
+      do_ref = base_do.clone()
+      Tensor.realize(do_ref)
+
+    q_ref_, k_ref_, v_ref_ = q_ref.transpose(1, 2), k_ref.transpose(1, 2), v_ref.transpose(1, 2)
+    ref = flash_attention(q_ref_, k_ref_, v_ref_, is_causal=True)
+    ref = ref.float().transpose(1, 2)
+    ref.backward(do_ref)
+    Tensor.realize(q_ref.grad, k_ref.grad, v_ref.grad)
+
+    assert_allclose(q.grad, q_ref.grad, atol=1e-5, rtol=1e-5)
+    assert_allclose(v.grad, v_ref.grad, atol=1e-5, rtol=1e-5)
+    assert_allclose(k.grad, k_ref.grad, atol=1e-5, rtol=1e-5)
+
 if __name__ == "__main__":
   unittest.main()

test/unit/test_assign.py

@@ -2,6 +2,7 @@
 import unittest
 import numpy as np
 from tinygrad import dtypes, Tensor, TinyJit, GlobalCounters, Variable
+from tinygrad.uop.ops import Ops
 from tinygrad.device import is_dtype_supported
 from tinygrad.helpers import temp, CI, CPU_LVP, Context
 
@@ -128,6 +129,7 @@ class TestAssign(unittest.TestCase):
     new = a + old_a
     np.testing.assert_allclose(new.numpy(), 4)
 
+  @unittest.skip("TODO: this is broken")
   def test_assign_changes_alt(self, realize=False):
     a = Tensor(1).contiguous()
     if realize: a.realize()
@@ -231,7 +233,6 @@ class TestAssign(unittest.TestCase):
     np.testing.assert_equal(b0.numpy(), 128)
     np.testing.assert_equal(b1.numpy(), 608)
 
-  @unittest.skip("TODO: bring this assert back")
   def test_crossunder_assign(self):
     # NOTE: should *not* raise AssertionError from numpy
     with self.assertRaisesRegex(RuntimeError, "cycle"):
@@ -268,6 +269,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)
@@ -484,10 +495,10 @@ class TestAssign(unittest.TestCase):
     np.testing.assert_allclose(c.numpy(), [4.0, 3.0, 3.0, 4.0])
 
   def test_assign_bitcast_different_size(self):
-    # different-size bitcast creates a new tensor, not a view, so assign doesn't modify the original
+    # assign to a shape-changing bitcast view (only works on DISK currently)
     a = Tensor([0]*8, dtype=dtypes.uint8).realize()
     a.bitcast(dtypes.int64).assign(Tensor([12345], dtype=dtypes.int64)).realize()
-    np.testing.assert_equal(a.numpy(), [0]*8)
+    np.testing.assert_equal(a.numpy(), [0]*8)  # TODO: should be [57, 48, 0, 0, 0, 0, 0, 0] (little-endian 12345)
 
   @unittest.skip("don't use output buffer, and mismatch dtype no longer supported")
   def test_cast_assignment(self):
@@ -637,6 +648,7 @@ class TestAssignOrdering(unittest.TestCase):
     self.assertEqual(r1.item(), 4)
     self.assertEqual(r2.item(), 8)
 
+  @unittest.skip("TODO: this is broken")
   def test_write_read_write_chain(self):
     """Write, read, write chain - middle read must complete before second write."""
     buf = Tensor.zeros(4).contiguous().realize()
@@ -790,5 +802,79 @@ class TestAssignOrdering(unittest.TestCase):
     buf[2:3].assign(Tensor.full((1,), 3.0))
     self.assertEqual(buf.sum().realize().item(), 6.0)
 
+# TODO: assigns into views of unrealized non-BUFFER bases are silently dropped
+class TestAssignToUnrealizedView(unittest.TestCase):
+  def test_copy(self):
+    t = Tensor.zeros(2,2, dtype=dtypes.int).to("CPU:0").contiguous().realize()
+    c = t.to("CPU:1")  # unrealized COPY
+    self.assertIs(c.uop.base.op, Ops.COPY)
+    c[:, 1:2].assign(Tensor.ones(2,1, dtype=dtypes.int).to("CPU:1").contiguous().realize())
+    # TODO: should be [[0,1],[0,1]]
+    self.assertEqual(c.tolist(), [[0,0],[0,0]])
+
+  def test_contiguous(self):
+    t = Tensor([[1,2],[3,4]]).contiguous().realize()
+    c = t.permute(1,0).contiguous()  # unrealized CONTIGUOUS
+    self.assertIs(c.uop.base.op, Ops.CONTIGUOUS)
+    c[:, 1:2].assign(Tensor.ones(2,1, dtype=dtypes.int).contiguous().realize())
+    # TODO: should be [[1,1],[2,1]]
+    self.assertEqual(c.tolist(), [[1,3],[2,4]])
+
+  def test_contiguous_backward(self):
+    t = Tensor([[1,2],[3,4]]).contiguous().realize()
+    cb = t.contiguous_backward()  # unrealized CONTIGUOUS_BACKWARD
+    self.assertIs(cb.uop.base.op, Ops.CONTIGUOUS_BACKWARD)
+    cb[:, 1:2].assign(Tensor.ones(2,1, dtype=dtypes.int).contiguous().realize())
+    # TODO: should be [[1,1],[3,1]]
+    self.assertEqual(cb.tolist(), [[1,2],[3,4]])
+
+  def test_detach_copy(self):
+    t = Tensor.zeros(2,2, dtype=dtypes.int).to("CPU:0").contiguous().realize()
+    d = t.to("CPU:1").detach()  # DETACH(unrealized COPY)
+    self.assertIs(d.uop.base.op, Ops.COPY)
+    d[:, 1:2].assign(Tensor.ones(2,1, dtype=dtypes.int).to("CPU:1").contiguous().realize())
+    # TODO: should be [[0,1],[0,1]]
+    self.assertEqual(d.tolist(), [[0,0],[0,0]])
+
+  def test_detach_contiguous(self):
+    t = Tensor([[1,2],[3,4]]).contiguous().realize()
+    d = t.permute(1,0).contiguous().detach()  # DETACH(unrealized CONTIGUOUS)
+    self.assertIs(d.uop.base.op, Ops.CONTIGUOUS)
+    d[:, 1:2].assign(Tensor.ones(2,1, dtype=dtypes.int).contiguous().realize())
+    # TODO: should be [[1,1],[2,1]]
+    self.assertEqual(d.tolist(), [[1,3],[2,4]])
+
+  def test_alu(self):
+    a = Tensor([1,2,3,4]).contiguous().realize()
+    b = Tensor([5,6,7,8]).contiguous().realize()
+    c = a + b  # unrealized ADD
+    self.assertIs(c.uop.base.op, Ops.ADD)
+    c[:2].assign(Tensor([99, 99]).realize())
+    # TODO: silently dropped, should be [99,99,10,12] or raise an error
+    self.assertEqual(c.tolist(), [6,8,10,12])
+
+  def test_reduce(self):
+    a = Tensor([[1,2],[3,4]]).contiguous().realize()
+    r = a.sum(axis=0)  # unrealized REDUCE_AXIS
+    self.assertIs(r.uop.base.op, Ops.REDUCE_AXIS)
+    r[:1].assign(Tensor([99]).realize())
+    # TODO: silently dropped, should be [99,6] or raise an error
+    self.assertEqual(r.tolist(), [4,6])
+
+  def test_cast(self):
+    a = Tensor([1,2,3,4]).contiguous().realize()
+    c = a.float()  # unrealized CAST
+    self.assertIs(c.uop.base.op, Ops.CAST)
+    c[:2].assign(Tensor([99, 99], dtype=dtypes.float).realize())
+    # TODO: silently dropped, should be [99,99,3,4] or raise an error
+    self.assertEqual(c.tolist(), [1,2,3,4])
+
+  def test_const(self):
+    c = Tensor(5).reshape(1, 1).expand(2, 2)
+    self.assertIs(c.uop.base.op, Ops.CONST)
+    c[:, 1:2].assign(Tensor.ones(2,1, dtype=dtypes.int).contiguous().realize())
+    # TODO: silently dropped, should be [[5,1],[5,1]] or raise an error
+    self.assertEqual(c.tolist(), [[5,5],[5,5]])
+
 if __name__ == "__main__":
   unittest.main()

test/unit/test_callify.py

@@ -0,0 +1,111 @@
+import unittest
+from tinygrad import Tensor, dtypes
+
+class TestCallify(unittest.TestCase):
+  def test_basic(self):
+    a = Tensor([1.,2,3])
+    b = Tensor([4.,5,6])
+    out = a + b
+    out.callify()
+    self.assertListEqual(out.tolist(), [5.0, 7.0, 9.0])
+
+  def test_const(self):
+    out = Tensor(2.0) + Tensor(3.0)
+    out.callify()
+    self.assertEqual(out.item(), 5.0)
+
+  def test_sum(self):
+    out = Tensor.ones(16).contiguous().sum()
+    out.callify()
+    self.assertEqual(out.item(), 16.0)
+
+  def test_multi_output(self):
+    a = Tensor([1.,2,3])
+    b = Tensor([4.,5,6])
+    c = a + b
+    d = a * b
+    c.callify(d)
+    self.assertListEqual(c.tolist(), [5.0, 7.0, 9.0])
+    self.assertListEqual(d.tolist(), [4.0, 10.0, 18.0])
+
+  def test_two_callify_independent(self):
+    a = Tensor([1.,2,3])
+    b = Tensor([4.,5,6])
+    c = a + b
+    c.callify()
+
+    d = Tensor([10.,20,30])
+    e = Tensor([1.,1,1])
+    f = d - e
+    f.callify()
+
+    self.assertListEqual(c.tolist(), [5.0, 7.0, 9.0])
+    self.assertListEqual(f.tolist(), [9.0, 19.0, 29.0])
+
+  def test_two_callify_shared_input(self):
+    a = Tensor([1.,2,3]).contiguous().realize()
+    b = a + 1
+    b.callify()
+    c = a * 2
+    c.callify()
+    self.assertListEqual(b.tolist(), [2.0, 3.0, 4.0])
+    self.assertListEqual(c.tolist(), [2.0, 4.0, 6.0])
+
+  def test_chained_callify(self):
+    a = Tensor([1.,2,3])
+    b = a + 1
+    b.callify()
+    b.realize()
+    c = b + 1
+    c.callify()
+    self.assertListEqual(c.tolist(), [3.0, 4.0, 5.0])
+
+  def test_gemm(self):
+    a = Tensor.ones(8, 8).contiguous()
+    b = Tensor.eye(8).contiguous()
+    out = a @ b
+    out.callify()
+    lst = out.tolist()
+    for y in range(8):
+      for x in range(8):
+        self.assertEqual(lst[y][x], 1.0)
+
+  def test_int_dtype(self):
+    a = Tensor([1,2,3], dtype=dtypes.int)
+    b = Tensor([4,5,6], dtype=dtypes.int)
+    out = a + b
+    out.callify()
+    self.assertListEqual(out.tolist(), [5, 7, 9])
+
+  def test_reduce(self):
+    out = Tensor([1.,2,3,4]).sum()
+    out.callify()
+    self.assertEqual(out.item(), 10.0)
+
+  def test_multiple_ops(self):
+    a = Tensor([1.,2,3])
+    b = Tensor([4.,5,6])
+    out = (a + b) * (a - b)
+    out.callify()
+    self.assertListEqual(out.tolist(), [-15.0, -21.0, -27.0])
+
+  def test_double_callify(self):
+    a = Tensor([1.,2,3])
+    b = Tensor([4.,5,6])
+    out = a + b
+    out.callify()
+    out.callify()
+    self.assertListEqual(out.tolist(), [5.0, 7.0, 9.0])
+
+  def test_double_callify_multi_output(self):
+    a = Tensor([1.,2,3])
+    b = Tensor([4.,5,6])
+    c = a + b
+    d = a * b
+    c.callify(d)
+    c.callify(d)
+    self.assertListEqual(c.tolist(), [5.0, 7.0, 9.0])
+    self.assertListEqual(d.tolist(), [4.0, 10.0, 18.0])
+
+if __name__ == "__main__":
+  unittest.main()

test/unit/test_disk_tensor.py

@@ -74,18 +74,18 @@ class TestRawDiskBuffer(unittest.TestCase):
     _test_bitcasted(t, dtypes.float32, 0.0)
     _test_bitcasted(t, dtypes.uint32, 0)
     # pi in float16 stored via int16
-    t.assign(Tensor.full((128, 64), 0x4248, dtype=dtypes.uint16).bitcast(dtypes.uint8)).realize()
+    t.bitcast(dtypes.uint16).assign(Tensor.full((128, 64), 0x4248, dtype=dtypes.uint16)).realize()
     _test_bitcasted(t, dtypes.float16, 3.140625)
     _test_bitcasted(t, dtypes.float32, 50.064727)
     _test_bitcasted(t, dtypes.uint16, 0x4248)
     _test_bitcasted(t, dtypes.uint32, 0x42484248)
     # pi in float32 stored via float32
-    t.assign(Tensor.full((128, 32), 3.1415927, dtype=dtypes.float32).bitcast(dtypes.uint8)).realize()
+    t.bitcast(dtypes.float32).assign(Tensor.full((128, 32), 3.1415927, dtype=dtypes.float32)).realize()
     _test_bitcasted(t, dtypes.float32, 3.1415927)
     _test_bitcasted(t, dtypes.uint32, 0x40490FDB)
     # doesn't suport normal cast
     with self.assertRaises(NotImplementedError):
-      Tensor.empty((4,), dtype=dtypes.int16, device=f"disk:{tmp}").cast(dtypes.float16).realize()
+      Tensor.empty((4,), dtype=dtypes.int16, device=f"disk:{tmp}").cast(dtypes.float16).to(None).realize()
 
     # Those two should be moved to test_dtype.py:test_shape_change_bitcast after bitcast works on non-disk
     with self.assertRaises(RuntimeError):
@@ -178,6 +178,13 @@ class TestSafetensors(TempDirTestCase):
     import json
     assert json.loads(dat[8:8+sz])['__metadata__']['hello'] == 'world'
 
+  def test_safe_save_only_copy(self):
+    from tinygrad.helpers import GlobalCounters
+    t = Tensor.rand(10, 10).realize()
+    GlobalCounters.reset()
+    safe_save({"t": t}, self.tmp("test_copy.safetensors"))
+    assert GlobalCounters.global_ops == 0, f"safe_save should have no compute, got {GlobalCounters.global_ops} ops"
+
   def test_save_all_dtypes(self):
     for dtype in dedup(DTYPES_DICT.values()):
       if dtype in [dtypes.bfloat16]: continue # not supported in numpy
@@ -264,18 +271,20 @@ class TestDiskTensor(TempDirTestCase):
   def test_strided_read(self):
     # test non-contiguous (strided) read - should read elements at indices 0, 2, 4
     dt = Tensor([0, 1, 2, 3, 4, 5]).to(f"disk:{self.tmp('dt_strided_read')}")
-    result = dt[::2].tolist()
-    # TODO: dt[::2] selects indices 0, 2, 4, so result should be [0, 2, 4]
-    # self.assertEqual(result, [0, 2, 4])
-    self.assertEqual(result, [0, 1, 2])  # wrong!
+    with self.assertRaises(RuntimeError):
+      result = dt[::2].tolist()
+      # TODO: dt[::2] selects indices 0, 2, 4, so result should be [0, 2, 4]
+      # self.assertEqual(result, [0, 2, 4])
+      self.assertEqual(result, [0, 1, 2])  # wrong!
 
   def test_permuted_read(self):
     # test non-contiguous (permuted) read - should read transposed
     dt = Tensor([[0, 1, 2], [3, 4, 5]]).to(f"disk:{self.tmp('dt_permuted_read')}")
-    result = dt.T.tolist()
-    # TODO: transpose should give [[0, 3], [1, 4], [2, 5]]
-    # self.assertEqual(result, [[0, 3], [1, 4], [2, 5]])
-    self.assertEqual(result, [[0, 1], [2, 3], [4, 5]])  # wrong!
+    with self.assertRaises(RuntimeError):
+      result = dt.T.tolist()
+      # TODO: transpose should give [[0, 3], [1, 4], [2, 5]]
+      # self.assertEqual(result, [[0, 3], [1, 4], [2, 5]])
+      self.assertEqual(result, [[0, 1], [2, 3], [4, 5]])  # wrong!
 
   def test_write_ones(self):
     out = Tensor.ones(10, 10, device="CPU").contiguous()
@@ -303,10 +312,11 @@ class TestDiskTensor(TempDirTestCase):
   def test_strided_setitem(self):
     # test non-contiguous (strided) setitem - should set elements at indices 0, 2, 4
     dt = Tensor([1, 2, 3, 4, 5, 6]).to(f"disk:{self.tmp('dt_strided_setitem')}")
-    dt[::2] = Tensor([10, 20, 30])
-    # TODO: dt[::2] selects indices 0, 2, 4, so result should be [10, 2, 20, 4, 30, 6]
-    # self.assertEqual(dt.tolist(), [10, 2, 20, 4, 30, 6])
-    self.assertEqual(dt.tolist(), [10, 20, 30, 4, 5, 6])  # wrong!
+    with self.assertRaises(RuntimeError):
+      dt[::2] = Tensor([10, 20, 30])
+      # TODO: dt[::2] selects indices 0, 2, 4, so result should be [10, 2, 20, 4, 30, 6]
+      # self.assertEqual(dt.tolist(), [10, 2, 20, 4, 30, 6])
+      self.assertEqual(dt.tolist(), [10, 20, 30, 4, 5, 6])  # wrong!
 
   def test_advanced_setitem_not_supported(self):
     dt = Tensor.arange(12).reshape(3, 4).to(f"disk:{self.tmp('dt_advanced_setitem')}")
@@ -354,15 +364,10 @@ class TestDiskTensor(TempDirTestCase):
 
   def test_assign_with_bitcast(self):
     # bitcast assign is used in safe_save for writing header length
-    # bitcast on source side works, bitcast on target side raises
     t = Tensor.empty(16, device=f"disk:{self.tmp('dt_assign_bitcast')}", dtype=dtypes.uint8)
-    # correct way: bitcast the source to match target dtype
-    t[0:8].assign(Tensor([12345], dtype=dtypes.int64, device="CPU").bitcast(dtypes.uint8))
+    t[0:8].bitcast(dtypes.int64).assign([12345])
     val = int.from_bytes(t[0:8].data(), 'little')
     self.assertEqual(val, 12345)
-    # bitcast on target with non-broadcastable dtype raises
-    with self.assertRaises(RuntimeError):
-      t[0:4].bitcast(dtypes.int32).assign(Tensor([12345], dtype=dtypes.int64))
 
   def test_assign_to_bitcast_view(self):
     # assign float values to a float32 view of a uint8 disk buffer (used by safe_save)

test/unit/test_function.py

@@ -0,0 +1,197 @@
+import numpy as np
+import unittest
+from tinygrad.function import function
+from tinygrad import Tensor
+from tinygrad.uop.ops import UOp
+
+class TestFunction(unittest.TestCase):
+  def test_simple(self):
+    @function
+    def f(a:Tensor, b:Tensor) -> Tensor: return a+b
+
+    a = Tensor([1,2,3])
+    b = Tensor([4,5,6])
+    np.testing.assert_equal(f(a,b).numpy(), [5,7,9])
+
+  def test_simple_same(self):
+    @function
+    def f(a:Tensor, b:Tensor) -> Tensor: return a+b
+
+    a = Tensor([1,2,3])
+    np.testing.assert_equal(f(a,a).numpy(), [2,4,6])
+
+  def test_implicit(self):
+    inp = Tensor([7,8,9])
+    @function
+    def f(a:Tensor, b:Tensor) -> Tensor: return a+b+inp
+
+    a = Tensor([1,2,3])
+    b = Tensor([4,5,6])
+    np.testing.assert_equal(f(a,b).numpy(), [12,15,18])
+
+  def test_implicit_same_as_input(self):
+    inp = Tensor([7,8,9])
+    @function
+    def f(a:Tensor, b:Tensor) -> Tensor: return a+b+inp
+
+    a = Tensor([1,2,3])
+    np.testing.assert_equal(f(a, inp).numpy(), [15,18,21])
+
+  def test_implicit_2(self):
+    inp = Tensor([7,8,9])
+    @function
+    def f(a:Tensor, b:Tensor) -> Tensor:
+      return a+b+inp
+    inp2 = Tensor([7,8,10])
+    @function
+    def g(a:Tensor, b:Tensor) -> Tensor:
+      return a+b+inp2
+
+    a = Tensor([1,2,3])
+    b = Tensor([4,5,6])
+    c = f(a,b)
+    d = g(a,b)
+    c.realize(d)
+    np.testing.assert_equal(c.numpy(), [12,15,18])
+    np.testing.assert_equal(d.numpy(), [12,15,19])
+
+  def test_implicit_unrealized(self):
+    inp = Tensor([1,2,3]) + Tensor([4,5,6])
+    @function
+    def f(a:Tensor) -> Tensor: return a + inp
+
+    np.testing.assert_equal(f(Tensor([10,20,30])).numpy(), [15,27,39])
+
+  def test_detach(self):
+    @function
+    def f(a:Tensor, b:Tensor) -> Tensor: return a.detach() + b
+
+    a = Tensor([1,2,3])
+    b = Tensor([4,5,6])
+    np.testing.assert_equal(f(a, b).numpy(), [5,7,9])
+
+  def test_method(self):
+    class Foo:
+      def __init__(self): self.w = Tensor([10,20,30])
+      @function
+      def __call__(self, x:Tensor) -> Tensor: return x + self.w
+
+    foo = Foo()
+    np.testing.assert_equal(foo(Tensor([1,2,3])).numpy(), [11,22,33])
+
+  def test_grad_gemm(self):
+    @function
+    def f(a:Tensor, b:Tensor) -> Tensor: return a @ b
+
+    a = Tensor([[1.,2.],[3.,4.]], requires_grad=True)
+    b = Tensor([[5.,6.],[7.,8.]], requires_grad=True)
+    (f(a, b).contiguous() * b).sum().backward()
+    Tensor.realize(a, b, a.grad, b.grad)
+    # L = sum((a@b) * b), dL/d(a@b) = b, dL/da = b @ b^T, dL/db = a^T @ b + (a@b)
+    na, nb = a.numpy(), b.numpy()
+    np.testing.assert_allclose(a.grad.numpy(), nb @ nb.T)
+    np.testing.assert_allclose(b.grad.numpy(), na.T @ nb + na @ nb)
+
+  def test_grad_implicit(self):
+    w = Tensor([1., 2., 3.], requires_grad=True)
+    w.realize() # TODO: this is required
+    @function
+    def f(x:Tensor) -> Tensor: return x * w
+
+    x = Tensor([4., 5., 6.])
+    f(x).sum().backward()
+    np.testing.assert_allclose(w.grad.numpy(), [4., 5., 6.])
+
+  def test_symbolic_index(self):
+    table = Tensor([10,20,30,40]).contiguous().realize()
+    @function
+    def f(x:Tensor, start_pos:int|UOp) -> Tensor:
+      return x + table[start_pos]
+
+    v = UOp.variable("start_pos", 0, 3)
+    np.testing.assert_equal(f(Tensor([1,2,3]), v.bind(0)).numpy(), [11,12,13])
+
+  def test_symbolic_shape_input(self):
+    table = Tensor([10,20,30,40]).contiguous().realize()
+    @function
+    def f(x:Tensor) -> Tensor: return x * 2
+    sz = UOp.variable("sz", 1, 3)
+    slic = table[:sz.bind(2)]
+    np.testing.assert_equal(f(slic)[:2].numpy(), [20,40])
+
+  def test_nested_calls(self):
+    w = Tensor([10., 20., 30.])
+    @function
+    def f(a:Tensor) -> Tensor: return a + w
+    @function
+    def g(a:Tensor) -> Tensor: return a * w
+
+    a = Tensor([1., 2., 3.])
+    np.testing.assert_allclose(g(f(a)).numpy(), [110., 440., 990.])
+
+  def test_name(self):
+    @function
+    def f(a:Tensor) -> Tensor: return a + 1
+    assert f(Tensor([1])).uop.arg.name.endswith("f")
+
+  def test_method_name(self):
+    class Foo:
+      @function
+      def __call__(self, x:Tensor) -> Tensor: return x + 1
+    assert Foo()(Tensor([1])).uop.arg.name.endswith("Foo.__call__")
+
+  def test_callable_instance(self):
+    class Foo:
+      def __init__(self): self.w = Tensor([10,20,30])
+      def __call__(self, x:Tensor) -> Tensor: return x + self.w
+    foo = Foo()
+    f = function(foo)
+    np.testing.assert_equal(f(Tensor([1,2,3])).numpy(), [11,22,33])
+    assert f(Tensor([1,2,3])).uop.arg.name.endswith("Foo")
+
+  def test_iadd(self):
+    @function
+    def f(x:Tensor) -> Tensor:
+      x += 1
+      return x
+
+    a = Tensor([1,2,3]).realize()
+    np.testing.assert_equal(f(a).numpy(), [2,3,4])
+    np.testing.assert_equal(a.numpy(), [3,4,5])  # TODO: should be [1,2,3]
+
+  def test_implicit_assign(self):
+    a = Tensor([1,2,3])
+    a += 1
+    c = Tensor([2,2,2]).contiguous()
+    @function
+    def f(b:Tensor) -> Tensor: return a+b+c
+    b = Tensor([10,20,30]).realize()
+    np.testing.assert_equal(f(b).numpy(), [14,25,36])
+
+  def test_assign_input(self):
+    @function
+    def f(a:Tensor, b:Tensor) -> Tensor:
+      a.assign(b+1)
+      return a
+
+    a = Tensor([1,2,3]).realize()
+    b = Tensor([10,20,30]).realize()
+    np.testing.assert_equal(f(a,b).numpy(), [11,21,31])
+    np.testing.assert_equal(a.numpy(), [11,21,31])  # TODO: should be [1,2,3]
+    np.testing.assert_equal(b.numpy(), [10,20,30])
+
+  @unittest.expectedFailure
+  def test_assign_slice(self):
+    @function
+    def f(a:Tensor, b:Tensor) -> Tensor:
+      a[1:] = b[1:]+1
+      return a
+
+    a = Tensor([1,2,3]).realize()
+    b = Tensor([10,20,30]).realize()
+    np.testing.assert_equal(f(a,b).numpy(), [1,21,31])
+    np.testing.assert_equal(a.numpy(), [1,2,3])
+    np.testing.assert_equal(b.numpy(), [10,20,30])
+
+if __name__ == '__main__':
+  unittest.main()

test/unit/test_realize_is_realize.py

@@ -28,7 +28,25 @@ class TestRealizeIsRealized(unittest.TestCase):
     t = Tensor.ones(8).contiguous().shard((d, d), axis=0).realize()
     assert all(u.is_realized for u in t.uop.src)
 
-  # TODO: these are not realized after .realize() because they stay as consts / don't allocate buffers
+  def test_empty(self):
+    t = Tensor.empty(4, 4).realize()
+    assert not t.uop.is_realized
+
+  def test_disk(self):
+    with tempfile.NamedTemporaryFile() as f:
+      f.write(b'\x00' * 16)
+      f.flush()
+      t = Tensor.empty(4, dtype=dtypes.float32, device=f"disk:{f.name}").realize()
+      assert not t.uop.is_realized
+
+  def test_assign(self):
+    t = Tensor([1, 2, 3])
+    t += 1
+    t.realize()
+    assert t.uop.is_realized
+
+  # TODO: these are not realized after .realize()
+
   def test_const_not_realized(self):
     t = Tensor(3.14).realize()
     assert not t.uop.is_realized
@@ -37,17 +55,6 @@ class TestRealizeIsRealized(unittest.TestCase):
     t = Tensor.ones(4, 4).realize()
     assert not t.uop.is_realized
 
-  def test_empty_not_realized(self):
-    t = Tensor.empty(4, 4).realize()
-    assert t.uop.is_realized
-
-  def test_disk(self):
-    with tempfile.NamedTemporaryFile() as f:
-      f.write(b'\x00' * 16)
-      f.flush()
-      t = Tensor.empty(4, dtype=dtypes.float32, device=f"disk:{f.name}").realize()
-      assert t.uop.is_realized
-
   def test_none_not_realized(self):
     t = Tensor(None).realize()
     assert not t.uop.is_realized

test/unit/test_setitem_schedule.py

@@ -36,7 +36,8 @@ class TestSetitemInto(unittest.TestCase):
     self.assertEqual(GlobalCounters.kernel_count, 0)
     t.realize()
     self.assertEqual(GlobalCounters.kernel_count, 1)
-    self.assertEqual(GlobalCounters.global_mem, 4)
+    # TODO: this can be just 4 if empty goes through is_realized setitem path
+    self.assertEqual(GlobalCounters.global_mem, 4*(3*2+1)) # 3 elements had +1, 1 is assigned directly
     t[1].realize()
     t.realize()
     self.assertEqual(GlobalCounters.kernel_count, 1)

test/unit/test_tinyfs.py

@@ -0,0 +1,81 @@
+import json, math, os, socketserver, threading, unittest
+import numpy as np
+from tinygrad import Tensor, dtypes
+from extra.tinyfs.fetch_file import hash_file, _python_hash_1mb
+
+_chunks: dict[bytes, bytes] = {}
+
+class _Handler(socketserver.StreamRequestHandler):
+  def handle(self):
+    while line := self.rfile.readline():
+      cmd = line.decode().strip()
+      if cmd == "INFO":
+        self.wfile.write(json.dumps({"node0": ["node0", f"127.0.0.1:{self.server.server_address[1]}"]}).encode() + b"\r\n")
+      elif cmd.startswith("STORE_IN"):
+        data = self.rfile.read(int(cmd.split()[1]))
+        hashes = bytearray()
+        for i in range(math.ceil(len(data) / Tensor.CHUNK_SIZE)):
+          chunk = data[i*Tensor.CHUNK_SIZE:(i+1)*Tensor.CHUNK_SIZE].ljust(Tensor.CHUNK_SIZE, b'\0')
+          h = _python_hash_1mb(chunk)
+          _chunks[h] = chunk
+          hashes.extend(h)
+        self.wfile.write(hashes)
+      elif cmd.startswith("LOAD_IN"):
+        hashes = self.rfile.read(int(cmd.split()[1]))
+        self.wfile.write(json.dumps(["node0"] * (len(hashes) // 16)).encode() + b"\r\n")
+      elif cmd.startswith("CHUNK_OUT"):
+        size = int(cmd.split()[1])
+        self.wfile.write(_chunks.get(self.rfile.read(16), bytes(size))[:size])
+      self.wfile.flush()
+
+# regressed in 55d3a5def "preallocate all realized buffers"
+class TestTinyFS(unittest.TestCase):
+  @classmethod
+  def setUpClass(cls):
+    _chunks.clear()
+    cls._server = socketserver.ThreadingTCPServer(('127.0.0.1', 0), _Handler)
+    cls._server.daemon_threads = True
+    threading.Thread(target=cls._server.serve_forever, daemon=True).start()
+    os.environ["TINYFS_ENDPOINT"] = f"127.0.0.1:{cls._server.server_address[1]}"
+
+  @classmethod
+  def tearDownClass(cls):
+    _chunks.clear()
+    os.environ.pop("TINYFS_ENDPOINT", None)
+    cls._server.shutdown()
+    cls._server.server_close()
+
+  def test_store(self):
+    h = Tensor([1.0, 2.0, 3.0, 4.0]).fs_store().realize()
+    self.assertEqual(h.shape, (16,))
+    self.assertEqual(h.dtype, dtypes.uint8)
+
+  def test_store_deterministic(self):
+    a = Tensor([1.0, 2.0, 3.0, 4.0]).fs_store().realize()
+    b = Tensor([1.0, 2.0, 3.0, 4.0]).fs_store().realize()
+    np.testing.assert_array_equal(a.numpy(), b.numpy())
+
+  def test_store_different_data(self):
+    a = Tensor([1.0, 2.0, 3.0, 4.0]).fs_store().realize()
+    b = Tensor([5.0, 6.0, 7.0, 8.0]).fs_store().realize()
+    self.assertNotEqual(a.tolist(), b.tolist())
+
+  def test_roundtrip_uint8(self):
+    arr = np.arange(256, dtype=np.uint8)
+    loaded = Tensor(arr).fs_store().realize().fs_load(len(arr)).to("CPU")
+    np.testing.assert_array_equal(loaded.numpy(), arr)
+
+  def test_roundtrip_multichunk_uint8(self):
+    arr = np.random.default_rng(42).integers(0, 256, size=Tensor.CHUNK_SIZE + 1024, dtype=np.uint8)
+    loaded = Tensor(arr).fs_store().realize().fs_load(len(arr)).to("CPU")
+    np.testing.assert_array_equal(loaded.numpy(), arr)
+
+  def test_hash_matches_python_impl(self):
+    arr = np.arange(256, dtype=np.uint8)
+    h = Tensor(arr).fs_store().realize()
+    # the hash from fs_store should match the pure-Python hash_file reference
+    padded = arr.tobytes().ljust(Tensor.CHUNK_SIZE, b'\0')
+    self.assertEqual(h.data().tobytes(), hash_file(padded))
+
+if __name__ == "__main__":
+  unittest.main()

tinygrad/__init__.py

@@ -4,6 +4,7 @@ if int(os.getenv("TYPED", "0")):
   install_import_hook(__name__)
 from tinygrad.tensor import Tensor                                    # noqa: F401
 from tinygrad.engine.jit import TinyJit                               # noqa: F401
+from tinygrad.function import function                                # noqa: F401
 from tinygrad.uop.ops import UOp
 Variable = UOp.variable
 from tinygrad.dtype import dtypes                                     # noqa: F401

tinygrad/apps/llm.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 import sys, argparse, typing, re, unicodedata, json, uuid, time, functools
-from tinygrad import Tensor, nn, UOp, TinyJit, getenv
+from tinygrad import Tensor, nn, UOp, TinyJit, getenv, function
 from tinygrad.helpers import partition, DEBUG, Timing, GlobalCounters, stderr_log, colored
 from tinygrad.viz.serve import TCPServerWithReuse, HTTPRequestHandler
 
@@ -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,11 +132,15 @@ class TransformerBlock:
     q = apply_rope(q, freqs_cis)
     k = apply_rope(k, freqs_cis)
 
-    if not hasattr(self, "cache_kv"):
-      self.cache_kv = Tensor.zeros(2, B, self.n_kv_heads, self.max_context, self.head_dim, dtype=k.dtype, device=k.device).contiguous().realize()
-    self.cache_kv[:, :, :, start_pos:start_pos+T, :].assign(Tensor.stack(k, v))
-    k = self.cache_kv[0, :, :, 0:start_pos+T, :]
-    v = self.cache_kv[1, :, :, 0:start_pos+T, :]
+    # 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
@@ -144,6 +149,7 @@ class TransformerBlock:
     attn = self.attn_output(attn)
     return x + attn
 
+  @function
   def _feed_forward(self, h: Tensor) -> Tensor:
     h_norm = self.ffn_norm(h)
     if hasattr(self, 'ffn_gate_exps'):
@@ -156,6 +162,9 @@ class TransformerBlock:
     return h + self.ffn_down(gated)
 
   def __call__(self, x: Tensor, start_pos: int|UOp):
+    if not hasattr(self, "cache_kv"):
+      # TODO: how is the dtype of this determined?
+      self.cache_kv = Tensor.zeros(2, x.shape[0], self.n_kv_heads, self.max_context, self.head_dim, device=x.device).contiguous().realize()
     return self._feed_forward(self._attention(x, start_pos)).contiguous()
 
 class Transformer:
@@ -340,6 +349,10 @@ if __name__ == "__main__":
   # 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()

tinygrad/codegen/late/devectorizer.py

@@ -5,7 +5,7 @@ from dataclasses import dataclass, field
 from tinygrad.dtype import dtypes, ImageDType, DType, AddrSpace, Invalid, PtrDType
 from tinygrad.uop.ops import UOp, Ops, UPat, PatternMatcher, GroupOp, identity_element
 from tinygrad.uop.symbolic import uop_given_valid, parse_valid, invalid_gate
-from tinygrad.helpers import getenv, flatten, AMX, prod, ceildiv, IMAGE
+from tinygrad.helpers import getenv, flatten, AMX, prod, IMAGE
 from tinygrad.renderer import Renderer
 
 # ***** image load valid simplification *****
@@ -187,9 +187,9 @@ def _do_image_fixup(dt:ImageDType, idx:UOp) -> tuple[UOp, UOp, int, int]:
   buf = idx.src[0]
   x, valid = idx.src[1].get_idx(), idx.src[1].get_valid()
   h, w = dt.shape[0], dt.shape[1]
-  if IMAGE == 1 and valid is not None and (tp:=dt.size // 4) // 64:
-    h, w = max(([(1, tp)] * (tp < 16384)) + [(tp//64//k, 64*k) for k in range(ceildiv(tp//64, 16384), min(tp//64, 256)+1) if (tp//64) % k == 0],
-               key=lambda hw: len(_drop_valid_stmts(valid, uop_given_valid(valid, UOp.vectorize((x//4)%hw[1], x//(4*hw[1]))), *hw)))
+  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)))
     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

tinygrad/codegen/opt/postrange.py

@@ -7,7 +7,7 @@ from tinygrad.uop.ops import axis_letters, axis_colors, axis_to_pos
 from tinygrad.device import Buffer
 from tinygrad.dtype import dtypes, ImageDType
 from tinygrad.helpers import colored, BEAM, getenv, DEBUG, to_function_name, NOOPT, argsort, round_up, prod, merge_dicts, get_single_element, flatten
-from tinygrad.helpers import ALLOW_TF32, count, Context, ceildiv
+from tinygrad.helpers import ALLOW_TF32, count, Context
 from tinygrad.codegen.opt import Opt, OptOps, KernelOptError, check
 from tinygrad.codegen.simplify import pm_flatten_range
 from tinygrad.renderer import Renderer
@@ -353,26 +353,17 @@ def apply_opts(ast:UOp, ren:Renderer) -> UOp:
       k = hand_coded_optimizations(k)
   return k.get_optimized_ast(name_override=ast.arg.name if ast.arg is not None and ast.arg.name != "test" else None)
 
-# max image width (pixels): 16384. max image size: 4 * 16384 ** 2
-def _image_shape(dt):
-  if dt.base not in (dtypes.half, dtypes.float) or isinstance(dt, ImageDType) or dt.size > 4*16384*16384 or dt.nbytes()%64 != 0: return None
-  if dt.size <= 4 * 16384: return (1, dt.size // 4, 4)
-  if (pxls:=dt.size // 4) % 64: return None
-  # verify that a valid format exists
-  try: return next((pxls // 64 // k, 64 * k, 4) for k in range(ceildiv(pxls // 64, 16384), min(pxls // 64, 256)+1))
-  except StopIteration: return None
-
 def make_image(pa, off, idx):
-  if (idx.tag is None or idx.tag) and (shape:=_image_shape(dt:=pa.dtype)):
-    new_idx = idx.replace(src=(pa.replace(dtype=(dtypes.imageh if dt.base==dtypes.half else dtypes.imagef)(shape, shape[1] * 4 * dt.itemsize)), off),
-                          dtype=dtypes.float if dt.base == dtypes.half else idx.dtype)
+  if not isinstance(dt:=pa.dtype, ImageDType) and (idx.tag is None or idx.tag) and (shapes:=ImageDType.valid_dims(dt)):
+    new_pa = pa.replace(dtype=(dtypes.imageh if dt.base==dtypes.half else dtypes.imagef)(shapes[0] + (4,), shapes[0][1] * 4 * dt.itemsize))
+    new_idx = idx.replace(src=(new_pa, off), dtype=dtypes.float if dt.base == dtypes.half else idx.dtype)
     return new_idx if idx.tag or dt.base == dtypes.float else new_idx.cast(dtypes.half)
 
 pm_make_images = PatternMatcher([
   # ensure we dont create an unfoldable image store
   (UPat(Ops.STORE, src=(UPat.var("idx"),), allow_any_len=True, name="st"), lambda idx,st:
    st.replace(src=(idx.rtag(is_image:=any(c.op is Ops.RANGE and (c.vmax+1)%4 == 0 for c in idx.src[1].get_idx().split_uop(Ops.ADD))),
-                   st.src[1].cast(dtypes.float if is_image and _image_shape(idx.src[0].dtype) else idx.dtype.base)))),
+                   st.src[1].cast(dtypes.float if is_image and ImageDType.valid_dims(idx.src[0].dtype) else idx.dtype.base)))),
   (UPat(Ops.INDEX, src=(UPat(Ops.PARAM, name="pa"), UPat.var("off")), name="idx"), make_image),
   # remove double cast from image loads / stores
   (UPat(Ops.INDEX, src=(UPat(Ops.PARAM, name="pa"),), allow_any_len=True, name="idx").cast(dtypes.half).cast(dtypes.float), lambda idx,pa:

tinygrad/dtype.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 from typing import Final, ClassVar, Callable, Literal
 import math, struct, ctypes, functools
 from dataclasses import dataclass, fields
-from tinygrad.helpers import getenv, prod, round_up, next_power2, OSX
+from tinygrad.helpers import ceildiv, getenv, prod, round_up, next_power2, OSX
 from enum import Enum, auto
 
 class ConstFloat(float):
@@ -121,13 +121,25 @@ class ImageDType(PtrDType):
     if self._pitch != -1: return self._pitch
     imgw, imgh, itemsize_log = self.shape[1], self.shape[0], int(math.log2(self.itemsize))
     if OSX: return round_up(imgw, 256) * 4 * self.itemsize
-    pitchalign = max(6, 11 - int(math.log2(imgh))) if imgh > 1 else 6
+    # needs to be IMAGE_PITCH_ALIGN=256 for AMD
+    min_pitchalign = int(math.log2(v)) if (v := getenv("IMAGE_PITCH_ALIGN", 0)) > 0 else 6
+    pitchalign = max(min_pitchalign, 11 - int(math.log2(imgh))) if imgh > 1 else min_pitchalign
     align_up = max(1, (8 // itemsize_log + 1) - imgh // 32) if pitchalign == 6 else (2 ** (pitchalign - itemsize_log - 2))
 
     granularity = 128 if self.itemsize == 4 else 256
     pitch_add = (1 << pitchalign) if min(next_power2(imgw), round_up(imgw, granularity)) - align_up + 1 <= imgw and imgw > granularity//2 else 0
     return round_up(imgw * 4 * self.itemsize, 1 << pitchalign) + pitch_add
 
+  # get list of (height, width) that do not require pitch padding
+  @staticmethod
+  def valid_dims(ptr:PtrDType) -> list[tuple[int,int]]:
+    ALIGN, MAXW = getenv("IMAGE_PITCH_ALIGN", 256 if OSX else 64), 16384
+    if ptr.base not in (dtypes.half, dtypes.float) or ptr.size > 4*MAXW*MAXW or (ptr.size if OSX else ptr.nbytes()) % ALIGN != 0: return []
+    if OSX and (ptr.size // 4) % ALIGN: return [] # OSX has stricter requirements for height=1 images
+    pxls: int = ptr.size // 4
+    return ([(1, pxls)] * (pxls < MAXW) + [(pxls//ALIGN//k, ALIGN*k) for k in range(ceildiv(pxls//ALIGN, MAXW), min(pxls//ALIGN, MAXW//ALIGN)+1)
+                                           if (pxls//ALIGN)%k == 0] if pxls//ALIGN else [])
+
 class dtypes:
   @staticmethod
   @functools.cache
@@ -283,70 +295,47 @@ def float_to_bf16(x):
   return struct.unpack('f', struct.pack('I', u))[0]
 
 # fp8-float conversions based on https://gitlab.com/nvidia/headers/cuda-individual/cudart/-/blob/main/cuda_fp8.hpp
+# (bias, sig_bits, mant_mask, min_denorm_half, ovf_threshold, max_norm, min_norm)
+_fp8_cfg = {
+  dtypes.fp8e4m3: (7, 4, 0x7, 0x3F50000000000000, 0x407D000000000000, 0x7E, 0x3F90000000000000),
+  dtypes.fp8e5m2: (15, 3, 0x3, 0x3EE0000000000000, 0x40EE000000000000-1, 0x7B, 0x3F10000000000000),
+}
+
 def float_to_fp8(x: float, dtype: DType) -> int:
   assert dtype in dtypes.fp8s, "Only for fp8s"
   # e4m3 don't support inf, return 0x7f(+NaN) and 0xff(-NaN) to match jax
   # NaN is unordered, can't compare with zero, use math.copysign to get sign
   if dtype == dtypes.fp8e4m3 and not math.isfinite(x): return 0x7f if math.copysign(1, x) > 0 else 0xff
-  if dtype == dtypes.fp8e5m2 and math.isinf(x): return 0x7c if math.copysign(1, x) > 0 else 0xfc
-  config = {
-      dtypes.fp8e4m3: {"EXP_BIAS": 7, "SIGNIFICAND_BITS": 4, "MANTISSA_MASK": 0x7, "MINDENORM_O2": 0x3F50000000000000,
-              "OVERFLOW_THRESHOLD": 0x407D000000000000, "MAXNORM": 0x7E, "MINNORM": 0x3F90000000000000, "INF_VALUE": 0x7F},
-      dtypes.fp8e5m2: {"EXP_BIAS": 15, "SIGNIFICAND_BITS": 3, "MANTISSA_MASK": 0x3, "MINDENORM_O2": 0x3EE0000000000000,
-              "OVERFLOW_THRESHOLD": 0x40EE000000000000 - 1, "MAXNORM": 0x7B, "MINNORM": 0x3F10000000000000, "INF_VALUE": 0x7E}
-  }[dtype]
+  if dtype == dtypes.fp8e5m2 and not math.isfinite(x): return (0 if math.copysign(1, x) > 0 else 0x80) | (0x7c if math.isinf(x) else 0x7f)
+  bias, sig_bits, mant_mask, min_denorm_half, ovf_threshold, max_norm, min_norm = _fp8_cfg[dtype]
   xbits, = struct.unpack('Q', struct.pack('d', x))
-  FP8_DP_HALF_ULP = 1 << (53 - config["SIGNIFICAND_BITS"] - 1)
-  sign = ((xbits >> 63) & 1) << 7
-  exp = (((xbits >> 52) & 0x7FF) - 1023 + config["EXP_BIAS"])
-  mantissa = (xbits >> (53 - config["SIGNIFICAND_BITS"])) & config["MANTISSA_MASK"]
-  absx = xbits & 0x7FFFFFFFFFFFFFFF
-
-  if absx <= config["MINDENORM_O2"]: res = 0
-  elif absx > 0x7FF0000000000000: res = 0x7F if dtype == dtypes.fp8e4m3 else 0x7E | mantissa
-  elif absx > config["OVERFLOW_THRESHOLD"]: res = config["MAXNORM"]
-  elif absx >= config["MINNORM"]:
-    res = ((exp << (config["SIGNIFICAND_BITS"] - 1)) | mantissa)
-    round_bits = xbits & ((FP8_DP_HALF_ULP << 1) - 1)
-    if (round_bits > FP8_DP_HALF_ULP) or (round_bits == FP8_DP_HALF_ULP and (mantissa & 1)): res = res + 1
+  half_ulp = 1 << (52 - sig_bits)
+  sign, exp, mantissa, absx = ((xbits>>63)&1)<<7, ((xbits>>52)&0x7FF)-1023+bias, (xbits>>(53-sig_bits))&mant_mask, xbits&0x7FFFFFFFFFFFFFFF
+  if absx <= min_denorm_half: res = 0
+  elif absx > ovf_threshold: res = max_norm
+  elif absx >= min_norm:
+    res, round_bits = (exp << (sig_bits - 1)) | mantissa, xbits & ((half_ulp << 1) - 1)
+    if round_bits > half_ulp or (round_bits == half_ulp and mantissa & 1): res += 1
   else:
     shift = 1 - exp
-    mantissa |= 1 << (config["SIGNIFICAND_BITS"] - 1)
-    res = (mantissa >> shift)
-    round_bits = (xbits | (1 << (53 - 1))) & ((FP8_DP_HALF_ULP << (shift + 1)) - 1)
-    if (round_bits > (FP8_DP_HALF_ULP << shift)) or (round_bits == (FP8_DP_HALF_ULP << shift) and (res & 1)):
-      res = res + 1
-
-  res |= sign
-  return int(res)
+    mantissa |= 1 << (sig_bits - 1)
+    res, half = mantissa >> shift, half_ulp << shift
+    round_bits = (xbits | (1 << 52)) & ((half << 1) - 1)
+    if round_bits > half or (round_bits == half and res & 1): res += 1
+  return int(res | sign)
 
 def fp8_to_float(x: int, dtype: DType) -> float:
   assert dtype in dtypes.fp8s, "Only for fp8s"
-  ur = x << 8
-
-  if dtype == dtypes.fp8e5m2 and (ur & 0x7FFF) > 0x7C00: ur = 0x7FFF
-  elif dtype == dtypes.fp8e4m3:
-    sign = ur & 0x8000
-    exponent = ((ur & 0x7800) >> 1) + 0x2000
-    mantissa = (ur & 0x0700) >> 1
-    absx = x & 0x7F
-    if absx == 0x7F: ur = 0x7FFF
-    elif exponent == 0x2000:
-      if mantissa != 0:
-        mantissa <<= 1
-        while (mantissa & 0x0400) == 0:
-          mantissa <<= 1
-          exponent -= 0x0400
-        mantissa &= 0x03FF
-      else:
-        exponent = 0
-      ur = (sign | exponent) | mantissa
-    else:
-      ur = (sign | exponent) | mantissa
-
-  half_bytes = struct.pack('<H', ur)
-  float32_val = struct.unpack('e', half_bytes)[0]
-  return float(float32_val)
+  if (x & 0x7F) == 0: return -0.0 if x & 0x80 else 0.0
+  bias, sig_bits, *_ = _fp8_cfg[dtype]
+  mant_bits, exp_bits = sig_bits - 1, 8 - sig_bits
+  exp_max, mant_max = (1 << exp_bits) - 1, (1 << mant_bits) - 1
+  sign, exp, mantissa = (x >> 7) & 1, (x >> mant_bits) & exp_max, x & mant_max
+  if exp == exp_max:
+    if dtype == dtypes.fp8e5m2: return math.copysign(math.nan if mantissa else math.inf, -1 if sign else 1)
+    if mantissa == mant_max: return math.nan
+  val = (mantissa / (mant_max + 1)) * 2 ** (1 - bias) if exp == 0 else (1 + mantissa / (mant_max + 1)) * 2 ** (exp - bias)
+  return -val if sign else val
 
 def storage_fmt_for_dtype(dtype:DType): return 'H' if dtype == dtypes.bfloat16 else 'B' if dtype in dtypes.fp8s else dtype.fmt
 

tinygrad/engine/allocations.py

@@ -0,0 +1,147 @@
+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
+
+@dataclass
+class AllocCtx:
+  uop_list: list[UOp] = field(default_factory=list)
+  buffer_map: dict[UOp, UOp] = field(default_factory=dict)
+  bases: set[UOp] = field(default_factory=set)
+  assigns: list[UOp] = field(default_factory=list)
+  replacements: list[UOp] = field(default_factory=list)
+
+def tag_uop(ctx:AllocCtx, x:UOp):
+  if x.tag is not None: return None
+  ctx.uop_list.append(x)
+  return x.replace(tag=(len(ctx.uop_list)-1,))
+
+def disk_copy_is_buffer(ctx:AllocCtx, u:UOp):
+  # copies to disk are replaced with the disk buffer
+  to_disk = isinstance(u._device, str) and u._device.startswith(("DISK", "TINYFS"))
+  if to_disk: ctx.buffer_map[u] = UOp.new_buffer(u.device, u.shard_size, u.dtype).reshape(u.max_shard_shape)
+  # all copies from disk/numpy are realized into a real buffer
+  from_creation = isinstance(u.src[0]._device, str) and any(u.src[0]._device.startswith(x) for x in ["NPY", "DISK", "PYTHON", "TINYFS"])
+  if from_creation: return tag_uop(ctx, u)
+
+def apply_after(ctx:AllocCtx, u:UOp):
+  ctx.buffer_map[u] = u.src[0]
+
+# CONTIGUOUS and ASSIGN + parents are the only nodes that get updated
+add_tags = PatternMatcher([
+  (UPat(Ops.COPY, name="u"), disk_copy_is_buffer),
+  # no tag on copies that are assigned
+  (UPat(Ops.ASSIGN, src=(UPat(), UPat(Ops.COPY, name="c")), name="a"),
+   lambda a,c: a.replace(src=(a.src[0], c.rtag(())), tag=a.tag+c.tag) if a.tag and c.tag else None),
+  (UPat(Ops.AFTER, name="u"), apply_after),
+  (UPat({Ops.CONTIGUOUS, Ops.ASSIGN}, name="x"), tag_uop),
+  (UPat(GroupOp.All, name="x"), lambda ctx,x: tag_uop(ctx,x) if x in ctx.bases else None),
+])
+
+def replace_contig_with_assign(u:UOp):
+  # if size is 0, remove the contig
+  if u.size == 0: return u.src[0]
+  # no real contig for DISK/TINYFS tensors, they are left alone
+  if isinstance(u._device, str) and u._device.startswith(("DISK", "TINYFS")): return u.rtag(None)
+  dtype = u.dtype
+  if isinstance(dtype, ImageDType):
+    if prod(dtype.shape) != prod(u.max_shard_shape) or ([x for x in u.max_shard_shape if x != 1] or [1])[-1] % 4 != 0:
+      if DEBUG >= 1: print(f"demoting Image {dtype} with shape {u.max_shard_shape}")
+      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)
+
+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
+  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
+
+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),
+  # 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)),
+  # 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
+  (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):
+  if x.tag is None: return None
+  ret = x.replace(tag=None)
+  for t in x.tag:
+    original_uop: UOp = ctx.uop_list[t]
+    replace_uop = ret
+    while replace_uop.op is Ops.ASSIGN: replace_uop = replace_uop.src[0]
+    ctx.buffer_map[original_uop] = replace_uop.shrink_to(original_uop.shape)
+  ctx.assigns.append(ret)
+  return ret
+
+def append_after(ctx:AllocCtx, x:UOp):
+  ctx.assigns.append(x)
+
+def replace_input_buffer(ctx:AllocCtx, b:UOp):
+  ctx.replacements.append(b)
+  return UOp.param(len(ctx.replacements)-1, b.dtype, b.shape, b._device,
+                   b._min_max if b.op is Ops.BIND else None, b.src[0].arg[0] if b.op is Ops.BIND else None)
+
+pm_finalize_call = PatternMatcher([
+  (UPat(Ops.ASSIGN, name="x"), untag_and_append),
+  (UPat(Ops.AFTER, name="x"), append_after),
+  (UPat(Ops.COPY, name="x"), lambda ctx,x: append_after(ctx,x) if isinstance(x.device, str) and x.device.startswith(("DISK", "TINYFS")) else None),
+  # replace UNIQUE with LUNIQUE for CONST cache key normalization
+  (UPat(Ops.CONST, src=(UPat(Ops.UNIQUE), UPat(Ops.DEVICE, name="d")), name="b"), lambda b,d: b.replace(src=(d,))),
+])
+
+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),
+  # strip value from BIND for cache key normalization, so different values hit same cache
+  (UPat(Ops.BIND, src=(UPat(Ops.DEFINE_VAR), UPat(Ops.CONST)), name="b"), replace_input_buffer),
+])
+
+@track_rewrites(lambda _,ret: f"Process {pluralize('Buffer', len(ret[1]))}")
+def transform_to_call(big_sink:UOp) -> tuple[UOp, dict[UOp, UOp]]:
+  if VIZ: graph_rewrite(big_sink, PatternMatcher([]), name="View Tensor Graph")
+  # uop list is a list in the original_sink graph and we can map to the tags later
+  # here we build buffer map
+  dont_realize = {Ops.CONST, Ops.BUFFER, Ops.BIND, Ops.DEFINE_VAR, Ops.AFTER}
+  ctx = AllocCtx(bases=set([x.multibase for x in big_sink.src if x.base.op not in dont_realize]))
+
+  # this rewrite is "read-only", it adds simple things to buffer_map and may sink things on big_sink, bottom_up
+  # this is the only one where we have to be careful to not break the tensor graph
+  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")
+
+  # 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)
+  if VIZ: graph_rewrite(ret, PatternMatcher([]), name="View Call")
+  return ret, ctx.buffer_map

tinygrad/engine/jit.py

@@ -348,6 +348,8 @@ class TinyJit(Generic[ReturnType]):
         update_depends(depends, jit_cache)
         pruned, onetime = partition(jit_cache, lambda ei: any(b in depends for b in get_out_buffers_for_ei(ei)))
         if DEBUG >= 1: print(f"pruned from {len(jit_cache)} -> {len(pruned)} kernels")
+        # sync before re-executing onetime kernels
+        for dev in set(Device[b.device] for ei in onetime for b in ei.bufs if b is not None): dev.synchronize()
         # run the onetime kernels here
         for ei in onetime:
           for b in ei.bufs: cast(Buffer, b).ensure_allocated()

tinygrad/engine/schedule.py

@@ -1,10 +1,11 @@
-import time
+import time, inspect
 from typing import cast
 from collections import deque
-from tinygrad.uop.ops import UOp, Ops, buffers, UOpMetaClass, track_rewrites, PatternMatcher, UPat, graph_rewrite, graph_rewrite_map, gate_kernel_sink
+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, flatten, pluralize, SCACHE
+from tinygrad.helpers import DEBUG, cpu_profile, TracingKey, SPEC, pluralize, SCACHE, BASEDIR
 from tinygrad.engine.realize import ExecItem
 
 # **** schedule linearizer
@@ -14,7 +15,7 @@ def _unwrap_src(s: UOp) -> UOp:
   while len(s.src) and s.op not in {Ops.AFTER, Ops.BUFFER, Ops.PARAM, Ops.MSELECT, Ops.MSTACK, Ops.BIND}: s = s.src[0]
   return s
 
-def create_schedule(sched_sink:UOp) -> tuple[list[ExecItem], UOp]:
+def create_schedule(sched_sink:UOp) -> UOp:
   with cpu_profile(TracingKey("toposort sched_sink")):
     # build kernel dependency graph: edges from producer kernel to consumer kernels
     children: dict[UOp, list[UOp]] = {}
@@ -22,7 +23,7 @@ def create_schedule(sched_sink:UOp) -> tuple[list[ExecItem], UOp]:
     for u in sched_sink.toposort(gate_kernel_sink):
       if u.op is not Ops.AFTER: continue
       k = u.src[1]
-      assert k.op in {Ops.CALL, Ops.END}, f"AFTER src[1] should be KERNEL or END, not {k.op}"
+      assert k.op in {Ops.CALL, Ops.END, Ops.LINEAR}, f"AFTER src[1] should be CALL or END, not {k.op}"
       in_degree.setdefault(k, 0)
       if k.op is Ops.END: assert k.src[0].op is Ops.CALL, f"END src[0] should be KERNEL, not {k.src[0].op}"
       # WAR deps from rangeify are stored in AFTER src[2:]
@@ -46,142 +47,111 @@ def create_schedule(sched_sink:UOp) -> tuple[list[ExecItem], UOp]:
 
   with cpu_profile(TracingKey("linearize schedule")):
     queue: deque[UOp] = deque(k for k,v in in_degree.items() if v == 0)
-    pre_schedule: list[ExecItem] = []
-    buf_uops_list: list[UOp] = []
+    linearized: list[UOp] = []
     while len(queue):
       rk = queue.popleft()
-      k = rk.src[0] if rk.op is Ops.END else rk
-      assert k.op is Ops.CALL, f"unexpected op in queue: {k.op}"
-      buf_uops = tuple(_unwrap_src(s).buf_uop for s in k.src[1:] if s.op is not Ops.BIND)
-      pre_schedule.append(ExecItem(k.src[0], [], k.arg.metadata))
-      buf_uops_list.append(UOp.sink(*buf_uops))
+      if rk.op is Ops.LINEAR:
+        linearized.extend(rk.src)
+      else:
+        k = rk.src[0] if rk.op is Ops.END else rk
+        assert k.op is Ops.CALL, f"unexpected op in queue: {k.op}"
+        buf_uops = tuple(_unwrap_src(s).buf_uop for s in k.src[1:] if s.op is not Ops.BIND)
+        linearized.append(k.src[0].call(*buf_uops, metadata=k.arg.metadata))
       for x in children.get(rk, []):
         in_degree[x] -= 1
         if in_degree[x] == 0: queue.append(x)
+  return UOp(Ops.LINEAR, src=tuple(linearized))
 
-  return pre_schedule, UOp.sink(*buf_uops_list)
+def linear_to_schedule(linear:UOp) -> list[ExecItem]:
+  """Convert a LINEAR UOp to a list of ExecItems."""
+  schedule: list[ExecItem] = []
+  for si in linear.src:
+    ast, buf_uops = si.src[0], si.src[1:]
+    # create subbuffers if needed
+    if ast.op is Ops.BUFFER_VIEW:
+      base = buf_uops[1].buffer
+      assert isinstance(base, Buffer), "base can't be MultiBuffer"
+      buffers[buf_uops[0]] = base.view(buf_uops[0].arg, ast.dtype, ast.arg[1]*base.dtype.itemsize)
+    ubufs = [b.buffer for b in buf_uops]
+    metadata = si.arg.metadata
+    if any(isinstance(x, MultiBuffer) for x in ubufs):
+      assert all(isinstance(x, MultiBuffer) for x in ubufs), "kernel must all be multibuffer"
+      dnums = [x for x in ast.variables() if x.expr == '_device_num']
+      for j, bufs in enumerate(zip(*[x.bufs for x in cast(tuple[MultiBuffer, ...], ubufs)])):
+        schedule.append(ExecItem(ast, list(bufs), metadata, {dnums[0].expr:j} if len(dnums) else {}))
+    else:
+      schedule.append(ExecItem(ast, cast(list[Buffer|None], ubufs), metadata))
+  return schedule
 
 from tinygrad.engine.memory import memory_planner
-from tinygrad.schedule.rangeify import get_rangeify_map
-from tinygrad.schedule.multi import get_multi_map
+from tinygrad.schedule.rangeify import get_kernel_graph
+from tinygrad.uop.ops import PatternMatcher, UPat
 
-def replace_input_buffer(ctx:tuple[dict[UOp, UOp], dict[str, int], list[int], list[int]], b:UOp):
-  if (ret:=ctx[0].get(b, None)) is None:
-    # replace BUFFER with PARAM for cache key normalization (same as CALL)
-    ctx[0][b] = ret = UOp.param(ctx[2][0], b.dtype, b.shape, b.device)
-    ctx[2][0] += 1
-  return ret
-
-def replace_input_const(ctx:tuple[dict[UOp, UOp], dict[str, int], list[int], list[int]], b:UOp):
-  if (ret:=ctx[0].get(b, None)) is None:
-    # replace UNIQUE with LUNIQUE for CONST cache key normalization
-    ctx[0][b] = ret = b.replace(src=(UOp(Ops.LUNIQUE, arg=ctx[3][0]), b.src[1]))
-    ctx[3][0] += 1
-  return ret
-
-def strip_bind(ctx:tuple[dict[UOp, UOp], dict[str, int], list[int], list[int]], b:UOp):
-  var, val = b.src[0], b.src[1].arg
-  assert var.expr not in ctx[1] or ctx[1][var.expr] == val, f"bind mismatch on {var}, {ctx[1][var.expr]} != {val}"
-  ctx[1][var.expr] = val
-  return ctx[0].setdefault(b, b.replace(src=(b.src[0],)))
-
-pm_pre_sched_cache = 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 UNIQUE with LUNIQUE for CONST cache key normalization
-  (UPat(Ops.CONST, src=(UPat(Ops.UNIQUE), UPat(Ops.DEVICE)), name="b"), replace_input_const),
-  # 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"), strip_bind),
-])
-
-def create_new_buffer(ctx:dict[UOp, UOp], b:UOp):
-  if (ret:=ctx.get(b, None)) is None: ctx[b] = ret = UOp.new_buffer(b.device, b.arg, b.dtype)
+def create_new_buffer(ctx:tuple[dict[UOp, UOp], tuple[UOp, ...]], b:UOp):
+  if (ret:=ctx[0].get(b, None)) is None: ctx[0][b] = ret = UOp.new_buffer(b.device, b.arg, b.dtype)
   return ret
 
 pm_post_sched_cache = PatternMatcher([
+  (UPat(Ops.PARAM, name="x"), lambda ctx,x: ctx[1][x.arg].rtag() if x.tag is None else None),
   # create new BUFFERs for LUNIQUE BUFFERs from rangeify
   (UPat(Ops.BUFFER, src=(UPat(Ops.LUNIQUE), UPat(Ops.DEVICE)), name="b"), create_new_buffer),
-  # restore CONST back to original CONST
-  (UPat(Ops.CONST, src=(UPat(Ops.LUNIQUE), UPat(Ops.DEVICE)), name="b"), lambda ctx,b: ctx.get(b)),
-  # restore PARAM back to original BUFFER
-  (UPat(Ops.PARAM, src=(UPat(), UPat(Ops.DEVICE)), name="b"), lambda ctx,b: ctx.get(b)),
-  # restore BIND value stripped in pm_pre_sched_cache
-  (UPat(Ops.BIND, src=(UPat(Ops.DEFINE_VAR),), name="b"), lambda ctx,b: ctx.get(b)),
 ])
 
-schedule_cache: dict[bytes, tuple[list[ExecItem], UOp]] = {}
-@track_rewrites(lambda _,ret: f"Schedule {pluralize('Kernel', len(ret[1]))}")
-def complete_create_schedule_with_vars(big_sink:UOp) -> tuple[dict[UOp, UOp], list[ExecItem], dict[str, int]]:
-  # big_sink srcs are all the Tensors
+# the AFTER structure is already in LINEAR
+pm_collapse_after = PatternMatcher([
+  (UPat(Ops.AFTER, name="x"), lambda x: x.src[0])
+])
+
+schedule_cache: dict[bytes, UOp] = {}
+def lower_schedule_to_linear(big_sink:UOp) -> UOp|None:
   st = time.perf_counter()
-
-  # replace BUFFERs with PARAMs, CONSTs UNIQUE with LUNIQUE, strip BIND values for cache key, extract var_vals
-  input_buffers: dict[UOp, UOp] = {}
-  var_vals: dict[str, int] = {}
-  big_sink_cache = graph_rewrite(big_sink, pm_pre_sched_cache, ctx=(input_buffers, var_vals, [0], [0]), name="rewrite for sched cache")
-  sched_cache_key = big_sink_cache.key
-
-  if not SCACHE or (sc_ret:=schedule_cache.get(sched_cache_key, None)) is None:
-    # verify Tensors match the spec (on big_sink, we only need to do this if cache misses)
+  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)
-
-    # hack to preserve metadata
-    graph_rewrite_map(big_sink, pm_pre_sched_cache, ctx=({}, {}, [0], [0]), name="preserve metadata")
-
-    # tensor map is what we return
-    tensor_map: dict[UOp, UOp] = {}
-
-    if any(isinstance(x._device, tuple) for x in big_sink_cache.toposort()):
-      tensor_map |= get_multi_map(big_sink_cache)
-      big_sink_cache = big_sink_cache.substitute(tensor_map, name="Apply Multi Map")
-      big_sink_cache = UOp.sink(*flatten([x.src if x.op is Ops.MULTI else [x] for x in big_sink_cache.src]))
-
-    tensor_map |= get_rangeify_map(big_sink_cache)
-    big_sink = big_sink_cache.substitute(tensor_map, name="Apply Kernelize Map")
-
-    pre_schedule, buf_uops_sink = create_schedule(big_sink)
-
-    # save in schedule cache (include AFTERs in tensor_map so we don't need big_sink)
-    after_map = [(u, u.buf_uop) for u in big_sink.toposort() if u.op is Ops.AFTER]
-    tensor_map_sink = UOp.sink(*flatten([(k,v) for k,v in tensor_map.items()]), *flatten(after_map))
-    combined_sink = UOp.sink(tensor_map_sink, buf_uops_sink)
-    if SCACHE: schedule_cache[sched_cache_key] = (pre_schedule, combined_sink)
+    # 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
   else:
     # schedule cache hit
-    del big_sink_cache
-    pre_schedule, combined_sink = sc_ret
-
-  # replace all the PARAMs/LUNIQUEs back (single graph_rewrite for everything)
-  input_buffers_inverse = {v:k for k,v in input_buffers.items()}
-  combined = graph_rewrite(combined_sink, pm_post_sched_cache, ctx=input_buffers_inverse, name="unrewrite combined")
-  tensor_map_sink, buf_uops_sink = combined.src
-  tm_src = tensor_map_sink.src
-  tensor_map = {tm_src[i]:tm_src[i+1] for i in range(0, len(tm_src), 2)}
-
-  # add bufs to pre_schedule
-  schedule: list[ExecItem] = []
-  for i, si in enumerate(pre_schedule):
-    buf_uops = buf_uops_sink.src[i].src
-    # create subbuffers if needed
-    if si.ast.op is Ops.BUFFER_VIEW:
-      base = buf_uops[1].buffer
-      assert isinstance(base, Buffer), "base can't be MultiBuffer"
-      buffers[buf_uops[0]] = base.view(buf_uops[0].arg, si.ast.dtype, si.ast.arg[1]*base.dtype.itemsize)
-    ubufs = tuple(b.buffer for b in buf_uops)
-    if any(isinstance(x, MultiBuffer) for x in ubufs):
-      assert all(isinstance(x, MultiBuffer) for x in ubufs), "kernel must all be multibuffer"
-      dnums = [x for x in si.ast.variables() if x.expr == '_device_num']
-      for j, bufs in enumerate(zip(*[x.bufs for x in cast(tuple[MultiBuffer, ...], ubufs)])):
-        schedule.append(ExecItem(si.ast, list(bufs), si.metadata, si.fixedvars | ({dnums[0].expr:j} if len(dnums) else {})))
+    linear = sc_ret
+  if (DEBUG >= 1 and len(linear.src) > 1) or DEBUG >= 3:
+    for frm in inspect.stack():
+      if frm.filename == "<string>": continue
+      if frm.filename.startswith(str(BASEDIR / "apps")): break
+      if not frm.filename.startswith(str(BASEDIR)) and not frm.filename.endswith("/contextlib.py"): break
     else:
-      # ONE -> ONE
-      schedule.append(ExecItem(si.ast, list(ubufs), si.metadata, si.fixedvars))
+      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" | {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")
+
+pm_schedule = PatternMatcher([
+  (UPat(Ops.CALL, src=(UPat(Ops.SINK),), allow_any_len=True, name="big_sink"), lower_schedule_to_linear),
+])
+
+@track_rewrites(lambda _,ret: f"Schedule {pluralize('Kernel', len(ret[0]))}")
+def complete_create_schedule_with_vars(big_sink:UOp) -> tuple[list[ExecItem], dict[str, int]]:
+  # big_sink srcs are all the Tensors
+  linear = graph_rewrite(big_sink, pm_schedule, name="schedule to linear")
+
+  # vars used in the schedule
+  used_vars = set().union(*[{v.expr for v in si.src[0].variables()} for si in linear.src])
+  # get var_vals
+  var_vals: dict[str, int] = {}
+  for b in big_sink.src[1:]:
+    if b.op is Ops.BIND:
+      nm = b.src[0].expr
+      if nm not in used_vars: continue
+      val = b.src[1].arg
+      assert nm not in var_vals or var_vals[nm] == val, f"bind mismatch on {nm}, {var_vals[nm]} != {val}"
+      var_vals[nm] = val
+
+  # convert LINEAR to ExecItems
+  schedule: list[ExecItem] = linear_to_schedule(linear)
   with cpu_profile(TracingKey("memory planner")): schedule = memory_planner(schedule)
-
-  if (DEBUG >= 1 and len(schedule) > 1) or DEBUG >= 3:
-    print(f"scheduled {len(schedule):4d} kernels in {(time.perf_counter()-st)*1000:8.2f} ms"+\
-          f" | {' cache hit' if SCACHE and sc_ret is not None else 'CACHE MISS'} {sched_cache_key.hex()[:8]}"+\
-          f" | {len(UOpMetaClass.ucache)} uops in cache")
-
-  used_vars = set().union(*[{v.expr for v in si.ast.variables()} for si in schedule])
-  return tensor_map, schedule, {k:v for k,v in var_vals.items() if k in used_vars}
+  return schedule, var_vals

tinygrad/function.py

@@ -0,0 +1,62 @@
+import functools
+from typing import Generic, TypeVar, Callable, cast
+from tinygrad.helpers import Context, dedup, getenv
+from tinygrad.uop.ops import UOp, Ops, graph_rewrite, PatternMatcher, UPat
+from tinygrad.tensor import Tensor
+
+def add_to_ctx(ctx, x:UOp):
+  ret = x.param_like(len(ctx))
+  ctx.append(x)
+  return ret
+
+pm_ctx = PatternMatcher([
+  (UPat((Ops.BUFFER, Ops.BIND), name="x"), add_to_ctx),
+  (UPat((Ops.ASSIGN, Ops.CONTIGUOUS), name="x"),
+   lambda ctx,x: add_to_ctx(ctx,x) if not x.op_in_backward_slice_with_self(Ops.PARAM) else None),
+])
+
+ReturnType = TypeVar('ReturnType')
+class function(Generic[ReturnType]):
+  def __init__(self, fxn:Callable[..., ReturnType]):
+    self.fxn = fxn
+
+  def __get__(self, obj, objtype=None): return functools.partial(self.__call__, obj) if obj is not None else self
+
+  def __call__(self, *args, **kwargs) -> ReturnType:
+    input_uops: list[UOp] = [(t.uop if isinstance(t, Tensor) else t)
+                             for name,t in list(enumerate(args))+sorted(kwargs.items()) if isinstance(t, (Tensor, UOp))]
+
+    # use the base
+    #input_uops = [x.multibase for x in input_uops]
+
+    # deduplicate input_uops, keeping the first occurrence index for each unique uop
+    call_uops: list[UOp] = dedup(input_uops)
+
+    # disable realize/schedule while this is running
+    # run it and do surgery later
+    with Context(ALLOW_DEVICE_USAGE=getenv("DEVICE_IN_FUNCTION_BUG", 0)):
+      ret = self.fxn(*args, **kwargs)
+    assert isinstance(ret, Tensor), "only supports one tensor return for now"
+
+    # replace the known inputs with params (using deduplicated slots)
+    subs = {}
+    for i,x in enumerate(call_uops): subs[x] = x.param_like(i)
+    uret = ret.uop.substitute(subs)
+
+    # add contiguous to call_uops
+    #call_uops = [x.contiguous() for x in call_uops]
+
+    # the BUFFERs that are left are the implicit inputs
+    uret = graph_rewrite(uret, pm_ctx, call_uops, bottom_up=True, name="get_implicit_inputs")
+    name = getattr(self.fxn, '__qualname__', None) or type(self.fxn).__qualname__
+
+    # assign output
+    #pbuffer = uret.param_like(len(call_uops))
+    #assigned = pbuffer.assign(uret).sink()
+    #buffer = UOp.new_buffer(pbuffer.device, pbuffer.size, pbuffer.dtype).reshape(uret.shape)
+    #call = assigned.call(*call_uops, buffer, name=name)
+    #ret = buffer.after(call)
+
+    ret = uret.call(*call_uops, name=name)
+    return cast(ReturnType, Tensor(ret, device=ret.device))
+

tinygrad/gradient.py

@@ -13,14 +13,20 @@ def reduce_gradient(ctx:UOp, ret:UOp, op:Ops):
     return ((mask/broadcast_to_input(count)) * broadcast_to_input(ctx),)
   if op == Ops.MUL: return (broadcast_to_input(ctx * ret) / ret.src[0],)
 
-def call_gradient(ctx:UOp, k:UOp):
+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, set(params))
-  subst = dict(zip(params, args))
-  return (None,) + tuple(grads[p].substitute(subst) if p in grads else None for p in params)
+  grads = compute_gradient(fxn, ctx.param_like(len(args)), set(params))
+  ret: list[UOp|None] = [None]
+  for i,p in enumerate(params):
+    if p in grads:
+      # TODO: compact the args and remove unused ones
+      ret.append(grads[p].call(*args, ctx, name=(k.arg.name or "")+f"_backward_{i}"))
+    else:
+      ret.append(None)
+  return tuple(ret)
 
 # ctx is grad_output
 pm_gradient = PatternMatcher([

tinygrad/helpers.py

@@ -13,6 +13,7 @@ def prod(x:Iterable[T]) -> T|int: return functools.reduce(operator.mul, x, 1)
 OSX, WIN = platform.system() == "Darwin", sys.platform == "win32"
 CI = os.getenv("CI", "") != ""
 ARCH_X86 = any(x in platform.processor() for x in ("Intel", "i386", "x86_64"))
+BASEDIR = pathlib.Path(__file__).parent
 
 # fix colors on Windows, https://stackoverflow.com/questions/12492810/python-how-can-i-make-the-ansi-escape-codes-to-work-also-in-windows
 if WIN: os.system("")

tinygrad/nn/__init__.py

@@ -2,7 +2,6 @@ from __future__ import annotations
 import math
 from tinygrad.tensor import Tensor
 from tinygrad.dtype import dtypes
-from tinygrad.device import is_dtype_supported
 from tinygrad.helpers import prod, make_tuple, flatten, USE_ATOMICS
 from tinygrad.nn import optim, state, datasets  # noqa: F401
 
@@ -36,7 +35,7 @@ class BatchNorm:
     self.weight: Tensor|None = Tensor.ones(sz) if affine else None
     self.bias: Tensor|None = Tensor.zeros(sz) if affine else None
 
-    self.num_batches_tracked = Tensor.zeros(dtype='long' if is_dtype_supported(dtypes.long) else 'int', requires_grad=False)
+    self.num_batches_tracked = Tensor.zeros(dtype='long', requires_grad=False)
     if track_running_stats: self.running_mean, self.running_var = Tensor.zeros(sz, requires_grad=False), Tensor.ones(sz, requires_grad=False)
 
   def calc_stats(self, x:Tensor) -> tuple[Tensor, Tensor]:

tinygrad/nn/optim.py

@@ -8,7 +8,7 @@ class Optimizer:
   """
   Base class for all optimizers.
   """
-  def __init__(self, params: list[Tensor], lr: float, fused=FUSE_OPTIM):
+  def __init__(self, params: list[Tensor], lr: float, device=None, fused=FUSE_OPTIM):
     # if requires_grad is None, but being put into an optimizer, set it to True
     for x in params:
       if x.requires_grad is None: x.requires_grad_(True)
@@ -16,19 +16,18 @@ class Optimizer:
     self.params: list[Tensor] = dedup([x for x in params if x.requires_grad])
     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.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,
                      dtype=least_upper_dtype(dtypes.default_float, dtypes.float32))
     if self.fused: self.pos_params = list(itertools.accumulate(self.params, lambda x,y: x+y.numel(), initial=0))
 
-  @property
-  def device(self): return self.params[0].device
-
   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).contiguous()]
-    return [Tensor.zeros_like(t, dtype=param_dtype, requires_grad=False).contiguous() for t in self.params]
+    if self.fused: return [Tensor.zeros(self.pos_params[-1], dtype=param_dtype, device=self.device, requires_grad=False)]
+    if isinstance(self.device, tuple): return [Tensor.zeros_like(t, dtype=param_dtype, requires_grad=False) for t in self.params]
+    else: return [Tensor.zeros(t.shape, dtype=param_dtype, device=self.device, requires_grad=False) for t in self.params]
 
   def zero_grad(self):
     """
@@ -54,13 +53,14 @@ class Optimizer:
       # NOTE: contiguous is for speed
       out, extra = self._step([Tensor.cat(*[t.flatten() for t in self.params], dim=0)],
                               [Tensor.cat(*[unwrap(t.grad).contiguous().flatten() for t in self.params], dim=0)])
-      updated_params = [out[0][self.pos_params[i]:self.pos_params[i+1]].reshape(tt.shape) for i, tt in enumerate(self.params)]
+      updates = [out[0][self.pos_params[i]:self.pos_params[i+1]].reshape(tt.shape) for i, tt in enumerate(self.params)]
     else:
-      updated_params, extra = self._step(self.params, [unwrap(t.grad) for t in self.params])
-    for i, tt in enumerate(self.params): tt.assign(updated_params[i])
+      updates, extra = self._step(self.params, [unwrap(t.grad) for t in self.params])
+    for i, tt in enumerate(self.params): tt.assign(self._apply_update(tt, updates[i]))
     return extra+self.params+self.buffers
 
   def _step(self, params:list[Tensor], grads:list[Tensor]) -> tuple[list[Tensor], list[Tensor]]: raise NotImplementedError
+  def _apply_update(self, t:Tensor, up:Tensor) -> Tensor: return t.detach() - up.to(t.device)
 
 class OptimizerGroup(Optimizer):
   """
@@ -74,17 +74,17 @@ class OptimizerGroup(Optimizer):
   def schedule_step(self) -> list[Tensor]: return [x for o in self.optimizers for x in o.schedule_step()]
 
 # LARS is essentially just trust ratio to SGD so if we just set the trust coeff 0.0 it's just standard SGD.
-def SGD(params: list[Tensor], lr=0.001, momentum=0.0, weight_decay=0.0, nesterov=False, classic=False, fused=FUSE_OPTIM):
+def SGD(params: list[Tensor], lr=0.001, momentum=0.0, weight_decay=0.0, nesterov=False, classic=False, device=None, fused=FUSE_OPTIM):
   """
   Stochastic Gradient Descent (SGD) optimizer with optional momentum and weight decay.
 
   `classic` is a boolean flag that determines whether to use the popular momentum update rule or the classic momentum update rule.
   """
-  return LARS(params, lr, momentum, weight_decay, 0, None, nesterov, classic=classic, pre_wd=True, tcoef=0.0, fused=fused)
+  return LARS(params, lr, momentum, weight_decay, 0, None, nesterov, classic=classic, pre_wd=True, tcoef=0.0, device=device, fused=fused)
 
 # Muon applies the newton schulz algorithm on gradient. also can include momentum, nesterov, and weight decay
 def Muon(params: list[Tensor], lr=0.001, momentum=0.95, weight_decay=0.1, ns_steps=5, ns_coefficients=(3.4445, -4.775, 2.0315),
-         nesterov=True, fused=FUSE_OPTIM):
+         nesterov=True, device=None, fused=FUSE_OPTIM):
   """
   SGD with newton-schulz iteration and post momentum weight decay.
 
@@ -92,7 +92,8 @@ def Muon(params: list[Tensor], lr=0.001, momentum=0.95, weight_decay=0.1, ns_ste
   - Paper: https://arxiv.org/pdf/2502.16982
   """
   assert not fused, "FUSE_OPTIM not allowed for Muon optimizer"
-  return LARS(params, lr, momentum, weight_decay, ns_steps, ns_coefficients, nesterov, classic=False, pre_wd=False, tcoef=0.0, fused=fused)
+  return LARS(params, lr, momentum, weight_decay, ns_steps, ns_coefficients, nesterov,
+              classic=False, pre_wd=False, tcoef=0.0, device=None, fused=fused)
 
 class LARS(Optimizer):
   """
@@ -101,8 +102,8 @@ class LARS(Optimizer):
   - Paper: https://arxiv.org/abs/1708.03888v3
   """
   def __init__(self, params:list[Tensor], lr=0.001, momentum=0.9, weight_decay=1e-4, ns_steps=0, ns_coefficients=None,
-               nesterov=False, classic=True, pre_wd=True, tcoef=0.001, fused=FUSE_OPTIM):
-    super().__init__(params, lr, fused)
+               nesterov=False, classic=True, pre_wd=True, tcoef=0.001, device=None, fused=FUSE_OPTIM):
+    super().__init__(params, lr, device, fused)
     self.momentum, self.wd, self.ns_steps, self.ns_coefficients  = momentum, weight_decay, ns_steps, ns_coefficients
     self.nesterov, self.classic, self.pre_wd, self.tcoef = nesterov, classic, pre_wd, tcoef
     self.b = self._new_optim_param() if self.momentum else []
@@ -126,24 +127,24 @@ class LARS(Optimizer):
       if not self.pre_wd and self.wd > 0: t = t.detach() * (1.0 - self.wd * self.lr)
       # popular momentum does pre learning rate update
       if not self.classic: g = g * r * self.lr
-      ret.append((t.detach() - g).cast(t.dtype))
+      ret.append(g.cast(t.dtype))
     return ret, self.b
 
 # LAMB is essentially just the trust ratio part of LARS applied to Adam/W so if we just set the trust ratio to 1.0 it's just Adam/W.
-def AdamW(params: list[Tensor], lr=0.001, b1=0.9, b2=0.999, eps=1e-8, weight_decay=0.01, fused=FUSE_OPTIM):
+def AdamW(params: list[Tensor], lr=0.001, b1=0.9, b2=0.999, eps=1e-8, weight_decay=0.01, device=None, fused=FUSE_OPTIM):
   """
   AdamW optimizer with optional weight decay.
 
   - Paper: https://arxiv.org/abs/1711.05101v3
   """
-  return LAMB(params, lr, b1, b2, eps, weight_decay, adam=True, fused=fused)
-def Adam(params: list[Tensor], lr=0.001, b1=0.9, b2=0.999, eps=1e-8, fused=FUSE_OPTIM):
+  return LAMB(params, lr, b1, b2, eps, weight_decay, adam=True, device=device, fused=fused)
+def Adam(params: list[Tensor], lr=0.001, b1=0.9, b2=0.999, eps=1e-8, device=None, fused=FUSE_OPTIM):
   """
   Adam optimizer.
 
   - Paper: https://arxiv.org/abs/1412.6980
   """
-  return LAMB(params, lr, b1, b2, eps, 0.0, adam=True, fused=fused)
+  return LAMB(params, lr, b1, b2, eps, 0.0, adam=True, device=device, fused=fused)
 
 class LAMB(Optimizer):
   """
@@ -151,10 +152,10 @@ class LAMB(Optimizer):
 
   - Paper: https://arxiv.org/abs/1904.00962
   """
-  def __init__(self, params: list[Tensor], lr=0.001, b1=0.9, b2=0.999, eps=1e-6, weight_decay=0.0, adam=False, fused=FUSE_OPTIM):
-    super().__init__(params, lr, fused)
+  def __init__(self, params: list[Tensor], lr=0.001, b1=0.9, b2=0.999, eps=1e-6, weight_decay=0.0, adam=False, device=None, fused=FUSE_OPTIM):
+    super().__init__(params, lr, device, fused)
     self.b1, self.b2, self.eps, self.wd, self.adam = b1, b2, eps, weight_decay, adam
-    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()
 
@@ -175,5 +176,5 @@ class LAMB(Optimizer):
         r: Tensor|float = Tensor.where(r1 > 0, Tensor.where(r2 > 0, r1 / r2, 1.0), 1.0)
       else:
         r = 1.0
-      ret.append((t.detach() - self.lr * r * up).cast(t.dtype))
+      ret.append((self.lr * r * up).cast(t.dtype))
     return ret, [self.b1_t, self.b2_t] + self.m + self.v

tinygrad/nn/state.py

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

tinygrad/renderer/llvmir.py

@@ -97,9 +97,6 @@ base_rewrite = PatternMatcher([
                                                             f", {ldt(u.dtype)} {ctx[u]}, i32 {i}" for i,u in enumerate(x.src)])),
   # unary/binary/ternary ops
   (UPat(Ops.BITCAST, name="x"), lambda ctx,x: f"  {ctx[x]} = bitcast {ldt(x.src[0].dtype)} {ctx[x.src[0]]} to {ldt(x.dtype)}"),
-  # rewrite cast to bool to CMPNE 0
-  (UPat(Ops.CAST, name="x", dtype=dtypes.bool),
-   lambda ctx,x: f"  {ctx[x]} = {lop[x.src[0].dtype.scalar()][Ops.CMPNE]} {ldt(x.src[0].dtype)} {ctx[x.src[0]]}, zeroinitializer"),
   (UPat(Ops.CAST, name="x"), lambda ctx,x: f"  {ctx[x]} = {lcast(x.src[0].dtype, x.dtype)} {ldt(x.src[0].dtype)} {ctx[x.src[0]]} to {ldt(x.dtype)}"),
   (UPat(Ops.TRUNC, name="x"),
    lambda ctx,x: f"  {ctx[x]} = call {ldt(x.dtype)} @llvm.trunc.{ldt(x.dtype.scalar())}({ldt(x.src[0].dtype)} {ctx[x.src[0]]})"),

tinygrad/renderer/nir.py

@@ -26,7 +26,6 @@ aop = {**{x:u_aop for x in (dtypes.bool,)+dtypes.uints}, **{x:s_aop for x in dty
 def c(t:DType, u:bool=True) -> str: return "u" if t in dtypes.uints and u else ("i" if t in dtypes.ints else ("f" if t in dtypes.floats else "b"))
 def ncast(b:mesa.nir_builder, src:mesa.nir_def, it:DType, ot:DType) -> mesa.nir_def:
   if isinstance(it, PtrDType) and ot == dtypes.long: return src
-  if ot == dtypes.bool: return nalu(b, c(it, False)+'ne'+('u' if c(it) == 'f' else ''), src, nimm(b, 0, it))
   return nalu(b, f"{c(it)}2{c(it) if it in dtypes.ints and ot in dtypes.ints else c(ot, ot == dtypes.bool)}{ot.bitsize}", src)
 
 def nif(b:mesa.nir_builder, cond:mesa.nir_def, then_fn:Callable, else_fn:Callable):

tinygrad/renderer/ptx.py

@@ -28,7 +28,8 @@ asm_for_op: dict[Ops, Callable] = {
   Ops.OR: lambda d,a,b,dt, name: f"or.pred {d}, {a}, {b};" if dt == dtypes.bool else f"or.b{name[1:]} {d}, {a}, {b};",
   Ops.IDIV: lambda d,a,b,dt,name: f"div.{name} {d}, {a}, {b};", Ops.MOD: lambda d,a,b,dt,name: f"rem.{name} {d}, {a}, {b};",
   Ops.MAX: lambda d,a,b,dt,name: f"max.{name} {d}, {a}, {b};", Ops.CMPEQ: lambda d,a,b,dt,name: f"setp.eq.{name} {d}, {a}, {b};",
-  Ops.CMPLT: lambda d,a,b,dt,name: f"setp.lt.{name} {d}, {a}, {b};", Ops.CMPNE: lambda d,a,b,dt,name: f"setp.ne.{name} {d}, {a}, {b};",
+  Ops.CMPLT: lambda d,a,b,dt,name: f"setp.lt.{name} {d}, {a}, {b};",
+  Ops.CMPNE: lambda d,a,b,dt,name: f"setp.{'neu' if dtypes.is_float(dt) else 'ne'}.{name} {d}, {a}, {b};",
   Ops.MULACC: lambda d,a,b,c,dt,name: f"{'fma.rn' if dtypes.is_float(dt) else 'mad.lo'}.{name} {d}, {a}, {b}, {c};",
   Ops.WHERE: lambda d,a,b,c,dt,name: [f"@{a} mov.{name} {d}, {b};", f"@!{a} mov.{name} {d}, {c};"] if dt == dtypes.bool else \
     f"selp.{'b16' if name == 'f16' else name} {d}, {b}, {c}, {a};"
@@ -98,8 +99,6 @@ string_rewrite = PatternMatcher([
   (UPat(Ops.BITCAST, name="x", src=(UPat.var("a"),), allow_any_len=True), lambda ctx, x, a: f"mov.b{ctx.types[x.dtype][1:]} {ctx.r[x]}, {ctx.r[a]};"),
   (UPat(Ops.CAST, name="x", src=(UPat(dtype=dtypes.bool, name="a"),)),
    lambda ctx, x, a: f"selp.b{ctx.types[x.dtype][1:]} {ctx.r[x]}, {render_val(1, x.dtype)}, {render_val(0, x.dtype)}, {ctx.r[a]};"),
-  (UPat(Ops.CAST, name="x", dtype=dtypes.bool, src=(UPat.var("a"),)),
-   lambda ctx, x, a: f"setp.ne.b{ctx.types[a.dtype][1:]} {ctx.r[x]}, {ctx.r[a]}, {render_val(0, a.dtype)};"),
   (UPat(Ops.CAST, name="x", src=(UPat.var("a"),)),
    lambda ctx, x, a: f"cvt{modifier(x.dtype, a.dtype)}.{ctx.cast_types[x.dtype]}.{ctx.cast_types[a.dtype]} {ctx.r[x]}, {ctx.r[a]};"),
   # store / gated load / load

tinygrad/runtime/ops_amd.py

@@ -45,7 +45,7 @@ class AMDSignal(HCQSignal):
 
   def _sleep(self, time_spent_since_last_sleep_ms:int):
     # Reasonable to sleep for long workloads (which take more than 200ms) and only timeline signals.
-    if time_spent_since_last_sleep_ms > 200 and self.is_timeline and self.owner is not None: self.owner.iface.sleep(200)
+    if time_spent_since_last_sleep_ms > 200 and self.owner is not None: self.owner.iface.sleep(200)
 
 class AMDComputeQueue(HWQueue):
   def __init__(self, dev:AMDDevice):
@@ -605,7 +605,7 @@ class AMDProgram(HCQProgram):
       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)
       self.dev.allocator._copyout(pmc_buf:=memoryview(bytearray(self.dev.pmc_buffer.size)), self.dev.pmc_buffer)
-      Compiled.profile_events += [ProfilePMCEvent(self.dev.device, self.dev.prof_prg_counter, self.dev.pmc_sched, bytes(pmc_buf),
+      Compiled.profile_events += [ProfilePMCEvent(self.dev.device, self.prof_prg_counter, self.dev.pmc_sched, bytes(pmc_buf),
                                                   self.dev.prof_exec_counter)]
     if self.dev.sqtt_enabled:
       cast(AMDComputeQueue, self.dev.hw_compute_queue_t()).sqtt_stop(self.dev.sqtt_wptrs) \
@@ -625,7 +625,7 @@ class AMDProgram(HCQProgram):
 
         self.dev.allocator._copyout(sqtt_mv:=memoryview(bytearray(wptr)), buf)
         resbuf = (struct.pack('<Q', 0x11 | (4 << 13) | (0xf << 16) | (se << 24)) + bytes(sqtt_mv)) if self.dev.target[0] == 9 else bytes(sqtt_mv)
-        Compiled.profile_events += [ProfileSQTTEvent(self.dev.device, self.dev.prof_prg_counter, se, resbuf,
+        Compiled.profile_events += [ProfileSQTTEvent(self.dev.device, self.prof_prg_counter, se, resbuf,
                                                      bool((SQTT_ITRACE_SE_MASK.value >> se) & 1), self.dev.prof_exec_counter)]
     return res
 
@@ -865,23 +865,25 @@ class PCIIface(PCIIfaceBase):
     return AMDQueueDesc(ring=ring.cpu_view().view(fmt='I'), doorbell=self.dev_impl.doorbell64.view(doorbell_index * 8, 8, fmt='Q'), put_value=pv,
       read_ptr=gart.cpu_view().view(offset=rptr, size=8, fmt='Q'), write_ptr=gart.cpu_view().view(offset=wptr, size=8, fmt='Q'), params=rcvr_params)
 
-  def sleep(self, timeout):
-    if hasattr(self.pci_dev, 'irq_poller') and self.pci_dev.irq_poller is not None and (events_cnt:=len(self.pci_dev.irq_poller.poll(timeout))):
-      self.pci_dev.irq_fd.read(8 * events_cnt)
-    self.dev_impl.ih.interrupt_handler()
-    if self.dev_impl.is_err_state: raise RuntimeError("Device is in error state")
-
-  def on_device_hang(self):
+  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()
-    faults = [f for d in devs if (f:=d.iface.dev_impl.gmc.check_fault())]
     for d in devs:
-      if d.iface.dev_impl.recover():
+      d.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
         d.timeline_signal.value = d.timeline_value - 1
         d.error_state = None
-    raise RuntimeError(f"Device hang detected: {'; '.join(faults)}" if faults else "Device hang detected")
+
+  def sleep(self, timeout):
+    if hasattr(self.pci_dev, 'irq_poller') and self.pci_dev.irq_poller is not None and (events_cnt:=len(self.pci_dev.irq_poller.poll(timeout))):
+      self.pci_dev.irq_fd.read(8 * events_cnt)
+    self._collect_faults()
+    if self.dev_impl.is_err_state: raise RuntimeError("Device is in error state")
+
+  def on_device_hang(self):
+    self._collect_faults(reset=True)
+    raise RuntimeError("Device hang detected")
 
   def device_fini(self): self.dev_impl.fini()
 

tinygrad/runtime/ops_cl.py

@@ -61,8 +61,9 @@ class CLProgram:
         if isinstance(dt, ImageDType):
           fmt = cl.cl_image_format(cl.CL_RGBA, {2:cl.CL_HALF_FLOAT, 4:cl.CL_FLOAT}[dt.itemsize])
           desc = cl.cl_image_desc(cl.CL_MEM_OBJECT_IMAGE2D, dt.shape[1], dt.shape[0], image_row_pitch=dt.pitch, buffer=b)
-          b = checked(cl.clCreateImage(self.dev.context, cl.CL_MEM_READ_WRITE, fmt, desc, None, status:=ctypes.c_int32()), status)
-        check(cl.clSetKernelArg(self.kernel, real_i, ctypes.sizeof(b), ctypes.byref(b)))
+          img = checked(cl.clCreateImage(self.dev.context, cl.CL_MEM_READ_WRITE, fmt, desc, None, status:=ctypes.c_int32()), status)
+          check(cl.clSetKernelArg(self.kernel, real_i, ctypes.sizeof(img), ctypes.byref(img)))
+        else: check(cl.clSetKernelArg(self.kernel, real_i, ctypes.sizeof(b), ctypes.byref(b)))
     for i,v in enumerate(vals,start=i+1): check(cl.clSetKernelArg(self.kernel, i, 4, ctypes.byref(ctypes.c_int32(v))))
     if local_size is not None: global_size = cast(tuple[int,int,int], tuple(int(g*l) for g,l in zip(global_size, local_size)))
     event = cl.cl_event() if wait else None

tinygrad/runtime/ops_cpu.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 import platform, sys, ctypes, functools, time, mmap, threading, queue
-from tinygrad.helpers import to_mv, OSX, WIN, mv_address, wait_cond, suppress_finalizing, unwrap, data64_le
+from tinygrad.helpers import to_mv, OSX, WIN, mv_address, suppress_finalizing, unwrap, data64_le
 from tinygrad.helpers import CPU_CC, CPU_LVP, CPU_LLVM
 from tinygrad.device import BufferSpec, DMACPURef, CompilerSet
 from tinygrad.runtime.support.hcq import HCQCompiled, HCQAllocator, HCQBuffer, HWQueue, HCQArgsState, HCQSignal, HCQProgram, MMIOInterface
@@ -13,7 +13,9 @@ from tinygrad.uop.ops import sint
 
 class CPUSignal(HCQSignal):
   def _sleep(self, time_spent_since_last_sleep_ms:int):
-    if self.is_timeline and self.owner is not None: self.owner.tasks.join()
+    if self.is_timeline and self.owner is not None:
+      self.owner.tasks.join()
+      if self.owner.error_state is not None: raise self.owner.error_state
 
 class CPUWorker(threading.Thread):
   def __init__(self, dev, tasks, thread_id):
@@ -29,13 +31,15 @@ class CPUWorker(threading.Thread):
   def run(self):
     while True:
       cmd_iter = iter(self.tasks.get())
-      for cmd in cmd_iter:
-        threads, args_cnt = next(cmd_iter), next(cmd_iter)
-        args = [next(cmd_iter) for _ in range(args_cnt)]
-        for th in range(threads - 1): self.push_task(th, cmd, args)
-        cmd(self.thread_id, *args)
-        for th in range(threads - 1): self.pool[th].join()
-      self.tasks.task_done()
+      try:
+        for cmd in cmd_iter:
+          threads, args_cnt = next(cmd_iter), next(cmd_iter)
+          args = [next(cmd_iter) for _ in range(args_cnt)]
+          for th in range(threads - 1): self.push_task(th, cmd, args)
+          cmd(self.thread_id, *args)
+          for th in range(threads - 1): self.pool[th].join()
+      except Exception as e: self.dev.error_state = e
+      finally: self.tasks.task_done()
 
 class CPUComputeQueue(HWQueue):
   def _exec(self, tid, prg, bufs, *args):
@@ -43,7 +47,9 @@ class CPUComputeQueue(HWQueue):
     if 'core_id' in prg.runtimevars: vals[prg.runtimevars['core_id']] = tid
     prg.fxn(*map(ctypes.c_uint64, args[:bufs]), *map(ctypes.c_int64 if platform.machine() == "arm64" else ctypes.c_int32, vals))
   def _signal(self, tid, signal_addr, value): to_mv(signal_addr, 4).cast('I')[0] = value
-  def _wait(self, tid, signal_addr, value): wait_cond(lambda: to_mv(signal_addr, 4).cast('I')[0] >= value, timeout_ms=60000)
+  def _wait(self, tid, tmpl_sig, signal_addr, value):
+    tmpl_sig.base_buf = HCQBuffer(signal_addr, 16, view=MMIOInterface(signal_addr, 16))
+    tmpl_sig.wait(value)
   def _timestamp(self, tid, timestamp_addr): to_mv(timestamp_addr, 8).cast('Q')[0] = time.perf_counter_ns()
   def cmd(self, cmd, *args, threads=1):
     self.q(cmd, threads, len(args), *args)
@@ -55,7 +61,7 @@ class CPUComputeQueue(HWQueue):
       self.bind_args_state(args_state)
       return self.cmd(self._exec, prg, 1, args_state.buf.va_addr)
     return self.cmd(self._exec, prg, len(args_state.bufs), *[x.va_addr for x in args_state.bufs], *args_state.vals, threads=(global_size or (1,))[0])
-  def wait(self, signal, value=0): return self.cmd(self._wait, signal.value_addr, value)
+  def wait(self, signal, value=0): return self.cmd(self._wait, type(signal)(signal.base_buf, owner=signal.owner, virt=True), signal.value_addr, value)
   def timestamp(self, signal): return self.cmd(self._timestamp, signal.timestamp_addr)
   def signal(self, signal, value:sint=0): return self.cmd(self._signal, signal.value_addr, value)
   def _submit(self, dev): dev.tasks.put(self._q[:])

tinygrad/runtime/ops_metal.py

@@ -156,7 +156,7 @@ class MetalAllocator(LRUAllocator[MetalDevice]):
     return MetalBuffer(ret, size)
   @suppress_finalizing
   def _free(self, opaque:MetalBuffer, options):
-    if not options.external_ptr: opaque.buf.release
+    if not options.external_ptr: opaque.buf.release()
   def _transfer(self, dest:MetalBuffer, src:MetalBuffer, sz:int, src_dev:MetalDevice, dest_dev:MetalDevice):
     dest_dev.synchronize()
     src_command_buffer = src_dev.mtl_queue.commandBuffer().retained()

tinygrad/runtime/ops_nv.py

@@ -28,7 +28,7 @@ class ProfilePMAEvent(ProfileEvent): device:str; kern:str; blob:bytes; exec_tag:
 class NVSignal(HCQSignal):
   def _sleep(self, time_spent_since_last_sleep_ms:int):
     # Reasonable to sleep for long workloads (which take more than 200ms) and only timeline signals.
-    if time_spent_since_last_sleep_ms > 200 and self.is_timeline and self.owner is not None: self.owner.iface.sleep(200)
+    if time_spent_since_last_sleep_ms > 200 and self.owner is not None: self.owner.iface.sleep(200)
 
 def get_error_str(status): return f"{status}: {nv_gpu.nv_status_codes.get(status, 'Unknown error')}"
 

tinygrad/runtime/support/am/amdev.py

@@ -254,8 +254,8 @@ class AMDev(PCIDevImplBase):
     else: self.mmio[reg] = val
 
   def wreg_pair(self, reg_base:str, lo_suffix:str, hi_suffix:str, val:int, inst:int=0):
-    self.reg(f"{reg_base}{lo_suffix}").write(val & 0xffffffff, inst=inst)
-    self.reg(f"{reg_base}{hi_suffix}").write(val >> 32, inst=inst)
+    self.reg(f"{reg_base}{lo_suffix}").write(lo32(val), inst=inst)
+    self.reg(f"{reg_base}{hi_suffix}").write(hi32(val), inst=inst)
 
   def indirect_rreg(self, reg:int) -> int:
     self.reg("regBIF_BX_PF0_RSMU_INDEX").write(reg * 4)
@@ -268,9 +268,9 @@ class AMDev(PCIDevImplBase):
   def indirect_wreg_pcie(self, reg:int, val:int, aid:int=0):
     reg_addr = reg * 4 + ((((aid & 0b11) << 32) | (1 << 34)) if aid > 0 else 0)
     self.reg("regBIF_BX0_PCIE_INDEX2").write(lo32(reg_addr))
-    if reg_addr >> 32: self.reg("regBIF_BX0_PCIE_INDEX2_HI").write(hi32(reg_addr) & 0xff)
+    if hi32(reg_addr) > 0: self.reg("regBIF_BX0_PCIE_INDEX2_HI").write(hi32(reg_addr) & 0xff)
     self.reg("regBIF_BX0_PCIE_DATA2").write(val)
-    if reg_addr >> 32: self.reg("regBIF_BX0_PCIE_INDEX2_HI").write(0)
+    if hi32(reg_addr) > 0: self.reg("regBIF_BX0_PCIE_INDEX2_HI").write(0)
 
   def _read_vram(self, addr, size) -> bytes:
     assert addr % 4 == 0 and size % 4 == 0, f"Invalid address {addr:#x} or size {size:#x}"