guard c1.arg

.github/workflows/benchmark.yml

@@ -21,9 +21,6 @@ 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:
@@ -44,14 +41,22 @@ 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
@@ -338,7 +343,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=120 NV=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py
+      run: BENCHMARK_LOG=cifar_10steps_half ASSERT_MIN_STEP_TIME=110 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
@@ -510,7 +515,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=230 AMD=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py
+      run: BENCHMARK_LOG=cifar_10steps_half ASSERT_MIN_STEP_TIME=200 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
@@ -520,9 +525,8 @@ jobs:
       run: time BENCHMARK_LOG=cifar AMD=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
     - name: Run full CIFAR training steps w 6 GPUS
       run: time BENCHMARK_LOG=cifar_6gpu AMD=1 DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
-    # this needs to be mocked and testable on a local machine
-    #- name: Test full tinyfs load
-    #  run: TINYFS_ENDPOINT=10.0.52.11:6767 PYTHONPATH=. python extra/tinyfs/fetch_file.py --hash d734f5e3be9f1e9d863bfaa4fc6c1ef2 --len 175866113 --dest mapping.json --check
+    - name: Test full tinyfs load
+      run: TINYFS_ENDPOINT=10.0.52.11:6767 PYTHONPATH=. python extra/tinyfs/fetch_file.py --hash d734f5e3be9f1e9d863bfaa4fc6c1ef2 --len 175866113 --dest mapping.json --check
     - name: Run process replay tests
       run: cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && PYTHONPATH=. python3 process_replay.py
 

.github/workflows/test.yml

@@ -649,8 +649,10 @@ 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 SQTT profiling tests
-        run: PROFILE=1 SQTT=1 python3 -m pytest -n=auto test/amd/test_sqtt_profiler.py
+      - 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 AMD emulated tests on NULL backend
         env:
           AMD: 0
@@ -662,30 +664,6 @@ 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
@@ -824,8 +802,6 @@ 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,5 +66,3 @@ 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_kernel_graph
+::: tinygrad.schedule.rangeify.get_rangeify_map
     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].flatten().assign(x.tobytes())
-    Y[idx].flatten().assign(y.tobytes())
+    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()
 
     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].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())
+      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()
 
-    imgs[idx].flatten().assign(img.tobytes())
+    imgs[idx].contiguous().realize().uop.base.realized.as_memoryview(force_zero_copy=True)[:] = img.tobytes()
 
     queue_out.put(idx)
   queue_out.put(None)

examples/mlperf/model_train.py

@@ -1371,9 +1371,8 @@ def train_llama3():
       # prevents memory spike on device 0
       v.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)
+  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)
 
   # init grads
   for p in optim.params:

examples/mlperf/optim.py

@@ -1,21 +1,14 @@
 from tinygrad.tensor import Tensor
 from tinygrad.dtype import dtypes
-from tinygrad.nn.optim import Optimizer
+from tinygrad.nn.optim import LAMB
 from tinygrad.helpers import 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()
+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)
     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()
@@ -28,19 +21,4 @@ class GradAccClipAdamW(Optimizer):
       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)
-
-    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)
+    return super()._step(params, grads)

extra/sqtt/examples/generate_examples.py

@@ -1,4 +1,4 @@
-import os, subprocess, sys
+import os, subprocess
 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.py TestCustomKernel.test_empty",
-  "test/test_tiny.py TestTiny.test_plus",
-  "test/test_tiny.py TestTiny.test_gemm",
+  "test.backend.test_custom_kernel.TestCustomKernel.test_empty",
+  "test.test_tiny.TestTiny.test_plus",
+  "test.test_tiny.TestTiny.test_gemm",
 ]
 
 if __name__ == "__main__":
@@ -17,8 +17,7 @@ if __name__ == "__main__":
   (EXAMPLES_DIR/arch).mkdir(exist_ok=True)
   for test in EXAMPLES:
     for i in range(2):
-      # 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":"."})
+      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)
       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,8 +11,7 @@ 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)
-  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))
+  shape = tuple(s // len(ref.device) if i == ref.uop.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)
 
@@ -30,40 +29,34 @@ 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
-  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'}")
+  B_local = B // num_devices
+  if DEBUG >= 2: print(f"Flash Attention {B=} {B_local=} {N=} {H=} {H_KV=} {D=}")
 
   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=shard_axis)
-  l_vec = _sharded_empty((B, H, 1, N), xq, dtype=dtypes.float32, axis=shard_axis_t)
+  attn = _sharded_empty_like(xq, axis=0)
+  l_vec = _sharded_empty((B, H, 1, N), xq, axis=0, dtype=dtypes.float32)
 
   def grad(dou:UOp, _) -> tuple[None, None, UOp, UOp, UOp]:
     do = Tensor(dou, device=dou.device)
-    dq_in = _sharded_empty((B, H, N, D), xq, axis=shard_axis_t)
-    dq = _sharded_empty_like(xq, axis=shard_axis)
-    dk = _sharded_empty_like(xk, axis=shard_axis)
-    dv = _sharded_empty_like(xv, axis=shard_axis)
+    dq_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)
 
     # delta_vec = (do * attn).sum(-1, dtype=dtypes.float32).transpose(1, 2).unsqueeze(-2).detach()
-    delta_vec = _sharded_empty((B, H, 1, N), xq, dtype=dtypes.float32, axis=shard_axis_t)
-    delta_vec, dq_in = Tensor.custom_kernel(delta_vec, dq_in, attn, do, fxn=functools.partial(custom_fa_backward_pre, device=single_device, arch=arch, B=B_local, N=N, H=H_local, H_KV=H_KV_local, D=D))[:2]
+    delta_vec = _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]
 
-    dq_in, dk, dv = Tensor.custom_kernel(dq_in, dk, dv, do, xq, xk, xv, l_vec, delta_vec, fxn=functools.partial(custom_fa_backward, device=single_device, arch=arch, B=B_local, N=N, H=H_local, H_KV=H_KV_local, D=D))[:3]
+    dq_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]
 
     # unshuffle dq
-    dq = Tensor.custom_kernel(dq, dq_in, fxn=functools.partial(custom_fa_backward_post, device=single_device, arch=arch, B=B_local, N=N, H=H_local, H_KV=H_KV_local, D=D))[0]
+    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]
 
     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_local, H_KV=H_KV_local, 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, H_KV=H_KV, D=D), grad_fxn=grad)[:2]
 
   return attn.transpose(1, 2)
 

extra/torch_backend/example.py

@@ -23,8 +23,7 @@ if __name__ == "__main__":
 
   kernel_count = GlobalCounters.kernel_count
   assert kernel_count > 0, "No kernels, test failed"
-  # NOTE: this is 124 on torch 2.10.0
-  expected_kernels = 332
+  expected_kernels = 228
   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, 5)
+    self._check_kernel_count(fn, 6)
 
   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, 5)
+    self._check_kernel_count(fn, 7)
 
   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, 7)
+    self._check_kernel_count(fn, 6)
 
   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, 7)
+    self._check_kernel_count(fn, 6)
 
   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, 7)
+    self._check_kernel_count(fn, 6)
 
   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, 17)
+    self._check_kernel_count(fn, 16)
 
   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, 28)
+    self._check_kernel_count(fn, 33)
 
 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": [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],
+    "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],
   }
 
 class TestSQTTExamplesRDNA4(SQTTExamplesTestBase): target = "gfx1200"

test/amd/test_sqtt_profiler.py

@@ -1,94 +0,0 @@
-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,7 +67,6 @@ 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,6 +265,8 @@ 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
@@ -273,7 +275,11 @@ 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
-    schedule = result.schedule()
+
+    big_sink = result.uop.sink()
+    tensor_map = get_rangeify_map(big_sink)
+    sched_sink = big_sink.substitute(tensor_map)
+    schedule, _ = create_schedule(sched_sink)
 
     # Find the custom_addmul kernel position
     custom_idx = next((i for i, item in enumerate(schedule)

test/backend/test_dtype.py

@@ -150,16 +150,28 @@ 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):
-    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())
+    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)
 
   @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):
-    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())
+    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)
 
   @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,6 +332,7 @@ 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()
@@ -475,7 +476,7 @@ class TestJit(unittest.TestCase):
     b = f(Tensor([2.0]))
     assert abs((a - b).item()) > 0.5
 
-  def test_jit_init_empty(self):
+  def test_jit_init_with_empty_different_size(self):
     @TinyJit
     def f(x:Tensor) -> Tensor: return (x + 1).realize()
 
@@ -484,16 +485,10 @@ class TestJit(unittest.TestCase):
     # scalar const input is not allowed
     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)
+    # list input has different view structure than empty(1)
+    # but okay if it's realized
+    #with self.assertRaises(JitError):
+    #  f(Tensor([2.0])).item()
 
 @unittest.skip("Pending multioutput implementation #3607")
 class TestMultioutputJit(unittest.TestCase):

test/backend/test_multitensor.py

@@ -135,6 +135,34 @@ 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)
@@ -676,7 +704,7 @@ class TestMultiTensor(unittest.TestCase):
 
     # test no left join
     with self.assertRaises((AssertionError, ValueError)):
-      t0.reshape((26*15,7)).contiguous().schedule()
+      t0.reshape((26*15,7)).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
@@ -812,15 +840,13 @@ 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):
-    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())
+    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())
 
   @unittest.skip("RANGEIFY doesn't support multi const folding")
   def test_multi_const_folding(self):
@@ -869,18 +895,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))).contiguous()
+      a = t.shrink(((0, 3), (0, 8)))
       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))).contiguous()
+    a = t.shrink(((0, 2), (0, 8)))
     a.schedule()
     assert a.shape == (2, 8)
 
-    p = a.pad(((0, 6), (0, 0))).contiguous()
+    p = a.pad(((0, 6), (0, 0)))
     p.schedule()
     assert p.shape == (8, 8)
 

test/backend/test_ops.py

@@ -8,8 +8,7 @@ from tinygrad.tensor import _to_np_dtype
 from tinygrad.device import is_dtype_supported
 from tinygrad.renderer.nir import NIRRenderer
 
-TINY_BACKEND = getenv("TINY_BACKEND")
-if TINY_BACKEND:
+if getenv("TINY_BACKEND"):
   import tinygrad.nn.torch # noqa: F401 # pylint: disable=unused-import
   torch.set_default_device("tiny")
 
@@ -419,6 +418,7 @@ 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,6 +640,8 @@ 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)
@@ -758,7 +760,6 @@ 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)
@@ -1542,6 +1543,7 @@ 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

@@ -0,0 +1,19 @@
+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, gc
+import unittest, struct, contextlib, statistics, time, 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(PROFILE=1):
+  with Context(VIZ=1, PROFILE=1):
     yield profile_list
     for dev in devs: dev.synchronize()
     for dev in devs: dev._at_profile_finalize()
@@ -170,19 +170,30 @@ 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):
+      with cpu_profile("test_1", dev.device):
         test_fxn(err=False)
       with self.assertRaises(Exception):
-        with cpu_profile("test_2", dev):
+        with cpu_profile("test_2", dev.device):
           test_fxn(err=True)
 
-    range_events = [p for p in profile if isinstance(p, ProfileRangeEvent) and p.device == dev]
+    range_events = [p for p in profile if isinstance(p, ProfileRangeEvent)]
     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,9 +78,7 @@ 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, 1))
+    run_schedule(check_schedule(expr, 0))
     np.testing.assert_allclose(expr.numpy(), np.full((4,), 4.0))
 
   def test_div_collapse_const(self):
     a = Tensor.full((4,), 4.0).contiguous().realize()
     expr = a/a
-    run_schedule(check_schedule(expr, 1))
+    run_schedule(check_schedule(expr, 0))
     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, 4))  # TODO: this is high
+    run_schedule(check_schedule(p, 3))
     tiny_ret = p.numpy()
 
     P = np.ones((3, 3), dtype=np.float32)
@@ -841,9 +841,10 @@ class TestSchedule(unittest.TestCase):
   def test_cast_const_view(self):
     a = Tensor.ones((4, 4), dtype=dtypes.float32)
     casted_view = a.cast(dtypes.int32)
-    run_schedule(check_schedule(casted_view, 1))
+    run_schedule(check_schedule(casted_view, 0))
+    self.assertIsNone(casted_view.uop.base.realized)
     realized_const_view = casted_view.contiguous()
-    run_schedule(check_schedule(realized_const_view, 0))
+    run_schedule(check_schedule(realized_const_view, 1))
     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))
@@ -1036,7 +1037,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, 6)  # TODO: this is high
+    sched = check_schedule(base, 2)
     run_schedule(sched)
     expected = list(range(16))
     for i, v in zip([1,2,5,6], [99,99,99,99]): expected[i] = v
@@ -1235,7 +1236,8 @@ class TestView(unittest.TestCase):
     bv = b.pad(((0, 2),))[-2:]
     # this becomes a late a*0
     late_mul = a*bv
-    run_schedule(check_schedule(late_mul, 2))
+    run_schedule(check_schedule(late_mul, 0))
+    # NOTE: no longer checked
     # the arange doesn't realize
     #self.assertIsNone(b.uop.base.realized)
     # mul doesn't realize
@@ -1252,7 +1254,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], 3)
+    s = check_schedule([late_mul, other_child], 2)
     # the arange becomes a BUFFER
     self.assertIs(b.uop.base.op, Ops.BUFFER)
     # NOTE: no longer checked
@@ -1265,7 +1267,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, 1, filter_sink=False))
+    run_schedule(check_schedule(b, 0, filter_sink=False))
     assert b.item() == 4
 
   def test_one_hot_with_copy(self):
@@ -1275,14 +1277,14 @@ class TestCopyFolding(unittest.TestCase):
 
   def test_const_copy_multi(self):
     x = Tensor.ones(1, device="CPU").to_(["CPU", "CPU:1"]) * 2
-    run_schedule(check_schedule(x, 2, filter_sink=False))
+    run_schedule(check_schedule(x, 0, filter_sink=False))
     self.assertEqual(x.item(), 2.0)
 
   def test_late_const_copy_folding(self):
     a = Tensor.arange(3).realize()
     zeros = Tensor.zeros(3).realize()
     b = (a*zeros).to("CPU") + 1
-    run_schedule(check_schedule(b, 1, filter_sink=False))
+    run_schedule(check_schedule(b, 0, filter_sink=False))
     self.assertListEqual(b.tolist(), [1, 1, 1])
     self.assertEqual(b.device, "CPU")
 
@@ -1322,7 +1324,7 @@ class TestCopyFolding(unittest.TestCase):
     a = Tensor.ones(4, 4).contiguous().realize()
     # use copy_to_device to bypass Tensor.to() shortcircuit and force a real same-device COPY in the graph
     a.assign(Tensor(a.uop.copy_to_device(a.device), a.device))
-    run_schedule(check_schedule(a, 2, filter_sink=False))
+    run_schedule(check_schedule(a, 0, 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, 5)
+    assert_jit_cache_len(jf, 4)
 
   def test_cat_dim0(self):
     def f(a, b): return a.cat(b, dim=0).realize()

test/backend/test_transcendental.py

@@ -84,6 +84,7 @@ 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,12 +113,6 @@ 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"} or getenv("AMD_IFACE", "") == "PCI", "Can't handle async update on CPU/MOCKAM device")
+  @unittest.skipIf(Device.DEFAULT in {"CPU"}, "Can't handle async update on CPU 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 and getenv("AMD_IFACE", "") != "PCI", "Emulate this on MOCKGPU to check the path in CI")
+  @unittest.skipUnless(MOCKGPU, "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, BufferSpec
+from tinygrad.device import CompileError, Device
 if Device.DEFAULT=="METAL":
   from tinygrad.runtime.ops_metal import MetalDevice, MetalCompiler, MetalProgram
 @unittest.skipIf(Device.DEFAULT!="METAL", "Metal support required")
@@ -48,14 +48,4 @@ 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)
-
-  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=}")
+      MetalProgram(device, "r_5", compiled)

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, functools, base64, codecs
+import os, multiprocessing, logging, pickle, sqlite3, difflib, warnings, itertools, functools, base64, codecs
 from dataclasses import replace
 from typing import Callable, Any
 
@@ -8,6 +8,7 @@ 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
@@ -42,6 +43,14 @@ 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()
@@ -59,6 +68,8 @@ 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, f"expected {expected_len}, got {len(fxn.jit_cache[0].prg.jit_cache)}"
+    assert len(fxn.jit_cache[0].prg.jit_cache) == expected_len
 
 def rand_for_dtype(dt:DType, size:int, allow_subnormal=True):
   if dtypes.is_unsigned(dt):

test/mockgpu/am/amdriver.py

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

@@ -1,314 +0,0 @@
-# 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(self.gpu.translate_addr(addr_lo + (addr_hi << 32)), 8)
+    ptr = to_mv(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(self.gpu.translate_addr(dst_addr_hi<<32|dst_addr_lo), 4).cast('I')[0] = self.gpu.regs[src_addr_lo]
+    to_mv(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(self.gpu.translate_addr(addr_lo + (addr_hi << 32)), 4).cast('I')[0]
+    elif mem_space == 1: mval = to_mv(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, self.gpu.translate_addr((addr_hi << 32) | addr_lo), buf_sz * 4, rptr, wptr)
+    self.ib_executor = PM4Executor(self.gpu, (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(self.gpu.translate_addr(struct.addr), 8).cast('Q')[0] = struct.data
+    to_mv(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(self.gpu.translate_addr(struct.addr), 4).cast('I')[0] & struct.mask
+    elif struct.mem_poll == 1: mval = to_mv(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(self.gpu.translate_addr(struct.addr), 8).cast('Q')
+    mem = to_mv(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] % 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)
+    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)
     self.rptr[0] += ctypes.sizeof(struct)
 
 class AMDGPURegisters:
@@ -343,7 +343,6 @@ 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,9 +1,7 @@
 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 ***
@@ -11,8 +9,7 @@ libc = ctypes.CDLL(ctypes.util.find_library("c"))
 libc.mmap.argtypes = [ctypes.c_void_p, ctypes.c_size_t, ctypes.c_int, ctypes.c_int, ctypes.c_int, ctypes.c_long]
 libc.mmap.restype = ctypes.c_void_p
 
-_amd_iface = getenv("AMD_IFACE", "")
-drivers = [NVDriver(), AMDriver() if _amd_iface == "PCI" else (AMUSBDriver() if _amd_iface == "USB" else AMDDriver())]
+drivers = [AMDDriver(), NVDriver()]
 tracked_fds = {}
 
 original_memoryview = builtins.memoryview
@@ -80,10 +77,9 @@ class MockFileIOInterface(FileIOInterface):
     return libc.mmap(start, sz, prot, flags, self.fd, offset)
 
   def read(self, size=None, 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].read_contents(size)
     if binary: raise NotImplementedError()
+    if self.fd in tracked_fds:
+      return tracked_fds[self.fd].read_contents(size)
     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)
@@ -93,20 +89,13 @@ class MockFileIOInterface(FileIOInterface):
       return tracked_fds[self.fd].list_contents()
     return os.listdir(self.path)
 
-  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 write(self, content, binary=False, offset=None): 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,213 +1,16 @@
-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_multitensor.py

@@ -38,27 +38,6 @@ 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 = {}
@@ -87,13 +66,5 @@ 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, 168, all_jitted=True)
+    helper_test("test_gpt2", lambda: (Tensor([[1,]]),Variable("pos", 1, 100).bind(1)), test, 0.23, 160, all_jitted=True)
 
   @slow
   def test_train_mnist(self):

test/null/test_rewrite_map.py

@@ -0,0 +1,202 @@
+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 not a.uop.is_realized
+    assert 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.is_realized
+    assert not a.uop.buffer.is_allocated()
 
   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, 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_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_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(), 13)
+      check_schedule(opt.schedule_step(), 19)
 
   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(), 13)
+      check_schedule(opt.schedule_step(), 19)
 
   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(), 15)
+      check_schedule(opt.schedule_step(), 21)
 
   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(), 11)
+      check_schedule(opt.schedule_step(), 13)
 
   def test_sgd_4convs_fuse(self):
     with Tensor.train():

test/null/test_tensor_metadata.py

@@ -4,7 +4,6 @@ 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,9 +390,6 @@ 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)")
 

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, 4)
+    self.assertLess(mem_ratio, 3)
 
   def test_dtype(self):
     IC, OC, X, Y = 4,4,9,9

test/testextra/test_cfg_viz.py

@@ -69,8 +69,6 @@ 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_dp(self):
+  def test_fast_fa_bwd_multidevice(self):
     Tensor.manual_seed(42)
 
     B, N, H, H_KV, D = 2, 1024, 32, 8, 128
@@ -175,52 +175,5 @@ 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,7 +2,6 @@
 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
 
@@ -129,7 +128,6 @@ 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()
@@ -233,6 +231,7 @@ 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"):
@@ -638,7 +637,6 @@ class TestAssignOrdering(unittest.TestCase):
     self.assertEqual(r1.item(), 4)
     self.assertEqual(r2.item(), 8)
 
-  @unittest.skip("TODO: this is broken")
   def test_write_read_write_chain(self):
     """Write, read, write chain - middle read must complete before second write."""
     buf = Tensor.zeros(4).contiguous().realize()
@@ -792,79 +790,5 @@ 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

@@ -1,120 +0,0 @@
-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_callify_then_schedule(self):
-    a = Tensor([1.,2,3])
-    b = Tensor([4.,5,6])
-    out = a + b
-    out.callify()
-    schedule = out.schedule()
-    self.assertGreater(len(schedule), 0)
-    self.assertListEqual(out.tolist(), [5.0, 7.0, 9.0])
-
-  def test_reduce(self):
-    out = Tensor([1.,2,3,4]).sum()
-    out.callify()
-    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

@@ -85,7 +85,7 @@ class TestRawDiskBuffer(unittest.TestCase):
     _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).to(None).realize()
+      Tensor.empty((4,), dtype=dtypes.int16, device=f"disk:{tmp}").cast(dtypes.float16).realize()
 
     # Those two should be moved to test_dtype.py:test_shape_change_bitcast after bitcast works on non-disk
     with self.assertRaises(RuntimeError):
@@ -264,20 +264,18 @@ 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')}")
-    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!
+    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')}")
-    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!
+    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()
@@ -305,11 +303,10 @@ 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')}")
-    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!
+    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')}")

test/unit/test_realize_is_realize.py

@@ -28,25 +28,7 @@ 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)
 
-  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()
-
+  # TODO: these are not realized after .realize() because they stay as consts / don't allocate buffers
   def test_const_not_realized(self):
     t = Tensor(3.14).realize()
     assert not t.uop.is_realized
@@ -55,6 +37,17 @@ 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,8 +36,7 @@ class TestSetitemInto(unittest.TestCase):
     self.assertEqual(GlobalCounters.kernel_count, 0)
     t.realize()
     self.assertEqual(GlobalCounters.kernel_count, 1)
-    # 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
+    self.assertEqual(GlobalCounters.global_mem, 4)
     t[1].realize()
     t.realize()
     self.assertEqual(GlobalCounters.kernel_count, 1)

test/unit/test_tinyfs.py

@@ -1,87 +0,0 @@
-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()
-
-  @unittest.expectedFailure
-  def test_store(self):
-    h = Tensor([1.0, 2.0, 3.0, 4.0]).fs_store().realize()
-    self.assertEqual(h.shape, (16,))
-    self.assertEqual(h.dtype, dtypes.uint8)
-
-  @unittest.expectedFailure
-  def test_store_deterministic(self):
-    a = Tensor([1.0, 2.0, 3.0, 4.0]).fs_store().realize()
-    b = Tensor([1.0, 2.0, 3.0, 4.0]).fs_store().realize()
-    np.testing.assert_array_equal(a.numpy(), b.numpy())
-
-  @unittest.expectedFailure
-  def test_store_different_data(self):
-    a = Tensor([1.0, 2.0, 3.0, 4.0]).fs_store().realize()
-    b = Tensor([5.0, 6.0, 7.0, 8.0]).fs_store().realize()
-    self.assertNotEqual(a.tolist(), b.tolist())
-
-  @unittest.expectedFailure
-  def test_roundtrip_uint8(self):
-    arr = np.arange(256, dtype=np.uint8)
-    loaded = Tensor(arr).fs_store().realize().fs_load(len(arr))
-    np.testing.assert_array_equal(loaded.numpy(), arr)
-
-  @unittest.expectedFailure
-  def test_roundtrip_multichunk_uint8(self):
-    arr = np.random.default_rng(42).integers(0, 256, size=Tensor.CHUNK_SIZE + 1024, dtype=np.uint8)
-    loaded = Tensor(arr).fs_store().realize().fs_load(len(arr))
-    np.testing.assert_array_equal(loaded.numpy(), arr)
-
-  @unittest.expectedFailure
-  def test_hash_matches_python_impl(self):
-    arr = np.arange(256, dtype=np.uint8)
-    h = Tensor(arr).fs_store().realize()
-    # 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/apps/llm.py

@@ -340,10 +340,6 @@ 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, IMAGE
+from tinygrad.helpers import getenv, flatten, AMX, prod, ceildiv, 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:
-    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)))
+  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)))
     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
+from tinygrad.helpers import ALLOW_TF32, count, Context, ceildiv
 from tinygrad.codegen.opt import Opt, OptOps, KernelOptError, check
 from tinygrad.codegen.simplify import pm_flatten_range
 from tinygrad.renderer import Renderer
@@ -353,17 +353,26 @@ 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 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)
+  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)
     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 ImageDType.valid_dims(idx.src[0].dtype) else idx.dtype.base)))),
+                   st.src[1].cast(dtypes.float if is_image and _image_shape(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 ceildiv, getenv, prod, round_up, next_power2, OSX
+from tinygrad.helpers import getenv, prod, round_up, next_power2, OSX
 from enum import Enum, auto
 
 class ConstFloat(float):
@@ -121,25 +121,13 @@ 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
-    # 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
+    pitchalign = max(6, 11 - int(math.log2(imgh))) if imgh > 1 else 6
     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
@@ -295,47 +283,70 @@ 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 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]
+  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]
   xbits, = struct.unpack('Q', struct.pack('d', x))
-  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
+  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
   else:
     shift = 1 - exp
-    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)
+    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)
 
 def fp8_to_float(x: int, dtype: DType) -> float:
   assert dtype in dtypes.fp8s, "Only for fp8s"
-  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
+  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)
 
 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

@@ -1,146 +0,0 @@
-from dataclasses import dataclass, field
-from tinygrad.uop.ops import UOp, UPat, PatternMatcher, Ops, GroupOp, graph_rewrite, identity_element, profile_matches
-from tinygrad.dtype import ImageDType
-from tinygrad.helpers import prod, DEBUG, argsort, VIZ
-
-@dataclass
-class AllocCtx:
-  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")
-  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"])
-  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 tensors, they are left alone
-  if isinstance(u._device, str) and u._device.startswith("DISK"): return u.rtag(None)
-  dtype = u.dtype
-  if isinstance(dtype, ImageDType):
-    if prod(dtype.shape) != prod(u.max_shard_shape) or ([x for x in u.max_shard_shape if x != 1] or [1])[-1] % 4 != 0:
-      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") 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),
-])
-
-@profile_matches
-def transform_to_call(big_sink:UOp) -> tuple[UOp, dict[UOp, UOp]]:
-  # 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,8 +348,6 @@ 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,10 @@
-import time, inspect
+import time
 from typing import cast
 from collections import deque
-from tinygrad.uop.ops import UOp, Ops, KernelInfo, buffers, UOpMetaClass, track_rewrites, graph_rewrite, gate_kernel_sink
+from tinygrad.uop.ops import UOp, Ops, buffers, UOpMetaClass, track_rewrites, PatternMatcher, UPat, graph_rewrite, graph_rewrite_map, gate_kernel_sink
 from tinygrad.uop.spec import type_verify, tensor_spec
 from tinygrad.device import Buffer, MultiBuffer
-from tinygrad.helpers import DEBUG, cpu_profile, TracingKey, SPEC, pluralize, SCACHE, BASEDIR
+from tinygrad.helpers import DEBUG, cpu_profile, TracingKey, SPEC, flatten, pluralize, SCACHE
 from tinygrad.engine.realize import ExecItem
 
 # **** schedule linearizer
@@ -14,7 +14,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) -> UOp:
+def create_schedule(sched_sink:UOp) -> tuple[list[ExecItem], UOp]:
   with cpu_profile(TracingKey("toposort sched_sink")):
     # build kernel dependency graph: edges from producer kernel to consumer kernels
     children: dict[UOp, list[UOp]] = {}
@@ -46,123 +46,142 @@ def create_schedule(sched_sink:UOp) -> UOp:
 
   with cpu_profile(TracingKey("linearize schedule")):
     queue: deque[UOp] = deque(k for k,v in in_degree.items() if v == 0)
-    linearized: list[UOp] = []
+    pre_schedule: list[ExecItem] = []
+    buf_uops_list: 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)
-      linearized.append(k.src[0].call(*buf_uops, metadata=k.arg.metadata))
+      pre_schedule.append(ExecItem(k.src[0], [], k.arg.metadata))
+      buf_uops_list.append(UOp.sink(*buf_uops))
       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)
 
 from tinygrad.engine.memory import memory_planner
-from tinygrad.schedule.rangeify import get_kernel_graph
-from tinygrad.uop.ops import PatternMatcher, UPat
+from tinygrad.schedule.rangeify import get_rangeify_map
+from tinygrad.schedule.multi import get_multi_map
 
-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)
+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)
   return ret
 
 pm_post_sched_cache = PatternMatcher([
-  # tag=True prevents re-matching after replacement (needed when PARAMs replace with PARAMs in nested callify)
-  (UPat(Ops.PARAM, name="x"), lambda ctx,x: ctx[1][x.arg].replace(tag=True) if x.tag is None else None),
   # create new BUFFERs for LUNIQUE BUFFERs from rangeify
   (UPat(Ops.BUFFER, src=(UPat(Ops.LUNIQUE), UPat(Ops.DEVICE)), name="b"), create_new_buffer),
+  # 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, UOp] = {}
-
-def _resolve_params(linear:UOp, params:tuple[UOp, ...]) -> UOp:
-  """Replace PARAMs in a LINEAR with the given params (BUFFERs or outer PARAMs), also handling LUNIQUE BUFFERs."""
-  from tinygrad.uop.ops import _remove_all_tags
-  linear = graph_rewrite(linear, pm_post_sched_cache, ctx=({}, params), name="params to buffers")
-  return graph_rewrite(linear, _remove_all_tags, name="remove tags")
-
-def rewrite_call_to_linear(ctx:list, call:UOp) -> UOp|None:
-  """Rewrite rule: CALL(SINK, *params) -> LINEAR(...) with caching. Only matches top-level CALLs from transform_to_call."""
-  function = call.src[0]
-  if function.op is not Ops.SINK or isinstance(function.arg, KernelInfo): return None
-  # recursively schedule any nested CALLs inside the function (from nested callify)
-  inner_start = len(ctx)
-  function = graph_rewrite(function, pm_schedule, ctx=ctx, name="schedule nested calls")
-  if not SCACHE or (linear:=schedule_cache.get(function.key, None)) is None:
-    if SPEC: type_verify(call.replace(src=(function,)+call.src[1:]), tensor_spec)
-    linear = create_schedule(get_kernel_graph(function))
-    if SCACHE: schedule_cache[function.key] = linear
-  # late apply params to buffers (tag=True prevents PARAM->PARAM cycles in nested callify)
-  linear = _resolve_params(linear, call.src[1:])
-  # resolve remaining PARAMs in inner LINEARs from nested CALLs using this call's params
-  for i in range(inner_start, len(ctx)):
-    inner_call, inner_linear = ctx[i]
-    ctx[i] = (inner_call, _resolve_params(inner_linear, call.src[1:]))
-  ctx.append((call, linear))
-  return linear
-
-pm_schedule = PatternMatcher([
-  (UPat(Ops.CALL, name="call"), rewrite_call_to_linear),
-  # strip AFTER(buf, LINEAR) -> buf after scheduling
-  (UPat(Ops.AFTER, src=(UPat(name="buf"), UPat(Ops.LINEAR))), lambda ctx,buf: buf),
-])
-
-def linear_to_schedule(linear:UOp) -> list[ExecItem]:
-  """Convert a LINEAR UOp to a list of ExecItems."""
-  schedule: list[ExecItem] = []
-  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, list(ubufs), metadata))
-  return schedule
-
-# strip AFTER(buf, LINEAR) -> buf, used by _apply_map_to_tensors to clean up scope tensors after scheduling
+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[list[UOp], list[ExecItem], dict[str, int]]:
+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
   st = time.perf_counter()
 
-  # rewrite CALLs to LINEARs and strip AFTERs
-  call_linear_pairs: list[tuple[UOp, UOp]] = []
-  graph_rewrite(big_sink, pm_schedule, ctx=call_linear_pairs, name="schedule calls")
-
-  # collect ExecItems from all LINEARs
-  schedule: list[ExecItem] = []
-  for _, linear in call_linear_pairs:
-    schedule.extend(linear_to_schedule(linear))
-
-  # get var_vals from CALL params
-  used_vars = set().union(*[{v.expr for v in si.src[0].variables()} for _, linear in call_linear_pairs for si in linear.src])
+  # 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] = {}
-  for call, _ in call_linear_pairs:
-    for b in call.src[1:]:
-      if b.op is Ops.BIND:
-        nm = b.src[0].expr
-        if nm not in used_vars: continue
-        val = b.src[1].arg
-        assert nm not in var_vals or var_vals[nm] == val, f"bind mismatch on {nm}, {var_vals[nm]} != {val}"
-        var_vals[nm] = val
+  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)
+    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)
+  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 {})))
+    else:
+      # ONE -> ONE
+      schedule.append(ExecItem(si.ast, list(ubufs), si.metadata, si.fixedvars))
   with cpu_profile(TracingKey("memory planner")): schedule = memory_planner(schedule)
 
   if (DEBUG >= 1 and len(schedule) > 1) or DEBUG >= 3:
-    for frm in inspect.stack():
-      if frm.filename.startswith(str(BASEDIR / "apps")): break
-      if not frm.filename.startswith(str(BASEDIR)) and not frm.filename.endswith("/contextlib.py"): break
-    else:
-      frm = None
-    print(f"scheduled {len(schedule):5d} kernels in {(time.perf_counter()-st)*1000:8.2f} ms"+\
-          f" | {len(UOpMetaClass.ucache):7d} uops in cache"+("" if frm is None else f" | {frm.filename}:{frm.lineno}"))
+    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")
 
-  return [call for call, _ in call_linear_pairs], schedule, var_vals
+  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}

tinygrad/helpers.py

@@ -13,7 +13,6 @@ 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/optim.py

@@ -8,7 +8,7 @@ class Optimizer:
   """
   Base class for all optimizers.
   """
-  def __init__(self, params: list[Tensor], lr: float, device=None, fused=FUSE_OPTIM):
+  def __init__(self, params: list[Tensor], lr: float, 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,18 +16,19 @@ 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)]
-    if isinstance(self.device, tuple): return [Tensor.zeros_like(t, dtype=param_dtype, requires_grad=False) for t in self.params]
-    else: return [Tensor.zeros(t.shape, dtype=param_dtype, device=self.device, requires_grad=False) for t in self.params]
+    if self.fused: return [Tensor.zeros(self.pos_params[-1], dtype=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]
 
   def zero_grad(self):
     """
@@ -53,14 +54,13 @@ 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)])
-      updates = [out[0][self.pos_params[i]:self.pos_params[i+1]].reshape(tt.shape) for i, tt in enumerate(self.params)]
+      updated_params = [out[0][self.pos_params[i]:self.pos_params[i+1]].reshape(tt.shape) for i, tt in enumerate(self.params)]
     else:
-      updates, extra = self._step(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]))
+      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])
     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, device=None, fused=FUSE_OPTIM):
+def SGD(params: list[Tensor], lr=0.001, momentum=0.0, weight_decay=0.0, nesterov=False, classic=False, 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, device=device, fused=fused)
+  return LARS(params, lr, momentum, weight_decay, 0, None, nesterov, classic=classic, pre_wd=True, tcoef=0.0, 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, device=None, fused=FUSE_OPTIM):
+         nesterov=True, fused=FUSE_OPTIM):
   """
   SGD with newton-schulz iteration and post momentum weight decay.
 
@@ -92,8 +92,7 @@ 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, device=None, fused=fused)
+  return LARS(params, lr, momentum, weight_decay, ns_steps, ns_coefficients, nesterov, classic=False, pre_wd=False, tcoef=0.0, fused=fused)
 
 class LARS(Optimizer):
   """
@@ -102,8 +101,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, device=None, fused=FUSE_OPTIM):
-    super().__init__(params, lr, device, fused)
+               nesterov=False, classic=True, pre_wd=True, tcoef=0.001, fused=FUSE_OPTIM):
+    super().__init__(params, lr, 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 []
@@ -127,24 +126,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(g.cast(t.dtype))
+      ret.append((t.detach() - 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, device=None, fused=FUSE_OPTIM):
+def AdamW(params: list[Tensor], lr=0.001, b1=0.9, b2=0.999, eps=1e-8, weight_decay=0.01, 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, 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):
+  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):
   """
   Adam optimizer.
 
   - Paper: https://arxiv.org/abs/1412.6980
   """
-  return LAMB(params, lr, b1, b2, eps, 0.0, adam=True, device=device, fused=fused)
+  return LAMB(params, lr, b1, b2, eps, 0.0, adam=True, fused=fused)
 
 class LAMB(Optimizer):
   """
@@ -152,10 +151,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, device=None, fused=FUSE_OPTIM):
-    super().__init__(params, lr, device, 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, fused=FUSE_OPTIM):
+    super().__init__(params, lr, 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) for _ in [b1, b2])
+    self.b1_t, self.b2_t = (Tensor.ones((1,), dtype=dtypes.float32, device=self.device, requires_grad=False).contiguous() for _ in [b1, b2])
     self.m = self._new_optim_param()
     self.v = self._new_optim_param()
 
@@ -176,5 +175,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((self.lr * r * up).cast(t.dtype))
+      ret.append((t.detach() - self.lr * r * up).cast(t.dtype))
     return ret, [self.b1_t, self.b2_t] + self.m + self.v

tinygrad/renderer/llvmir.py

@@ -97,6 +97,9 @@ 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,6 +26,7 @@ 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,8 +28,7 @@ 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.{'neu' if dtypes.is_float(dt) else '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.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};"
@@ -99,6 +98,8 @@ 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.owner is not None: self.owner.iface.sleep(200)
+    if time_spent_since_last_sleep_ms > 200 and self.is_timeline 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.prof_prg_counter, self.dev.pmc_sched, bytes(pmc_buf),
+      Compiled.profile_events += [ProfilePMCEvent(self.dev.device, self.dev.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.prof_prg_counter, se, resbuf,
+        Compiled.profile_events += [ProfileSQTTEvent(self.dev.device, self.dev.prof_prg_counter, se, resbuf,
                                                      bool((SQTT_ITRACE_SE_MASK.value >> se) & 1), self.dev.prof_exec_counter)]
     return res
 

tinygrad/runtime/ops_cl.py

@@ -61,9 +61,8 @@ 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)
-          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)))
+          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)))
     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, suppress_finalizing, unwrap, data64_le
+from tinygrad.helpers import to_mv, OSX, WIN, mv_address, wait_cond, 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,9 +13,7 @@ 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.owner.error_state is not None: raise self.owner.error_state
+    if self.is_timeline and self.owner is not None: self.owner.tasks.join()
 
 class CPUWorker(threading.Thread):
   def __init__(self, dev, tasks, thread_id):
@@ -31,15 +29,13 @@ class CPUWorker(threading.Thread):
   def run(self):
     while True:
       cmd_iter = iter(self.tasks.get())
-      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()
+      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()
 
 class CPUComputeQueue(HWQueue):
   def _exec(self, tid, prg, bufs, *args):
@@ -47,9 +43,7 @@ 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, tmpl_sig, signal_addr, value):
-    tmpl_sig.base_buf = HCQBuffer(signal_addr, 16, view=MMIOInterface(signal_addr, 16))
-    tmpl_sig.wait(value)
+  def _wait(self, tid, signal_addr, value): wait_cond(lambda: to_mv(signal_addr, 4).cast('I')[0] >= value, timeout_ms=60000)
   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)
@@ -61,7 +55,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, type(signal)(signal.base_buf, owner=signal.owner, virt=True), signal.value_addr, value)
+  def wait(self, signal, value=0): return self.cmd(self._wait, 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.owner is not None: self.owner.iface.sleep(200)
+    if time_spent_since_last_sleep_ms > 200 and self.is_timeline 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/hcq.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 from typing import cast, Callable, Type, TypeVar, Generic, Any
-import contextlib, decimal, statistics, time, ctypes, array, os, struct, collections, functools, itertools
+import contextlib, decimal, statistics, time, ctypes, array, os, struct, collections, functools
 try: import fcntl # windows misses that
 except ImportError: fcntl = None #type:ignore[assignment]
 from tinygrad.helpers import PROFILE, getenv, to_mv, from_mv, cpu_profile, ProfileRangeEvent, select_first_inited, unwrap, suppress_finalizing
@@ -214,26 +214,23 @@ class HWQueue(Generic[SignalType, HCQDeviceType, ProgramType, ArgsStateType]):
   def _submit(self, dev:HCQDeviceType): raise NotImplementedError("need _submit")
 
 class HCQSignal(Generic[HCQDeviceType]):
-  def __init__(self, base_buf:HCQBuffer, value:int=0, owner:HCQDeviceType|None=None, is_timeline:bool=False, timestamp_divider=1000, virt=False):
-    self.base_buf, self.owner, self.is_timeline = base_buf, owner, is_timeline
-    self.should_return = isinstance(self.base_buf.va_addr, int) and self.owner is not None and not virt
+  def __init__(self, base_buf:HCQBuffer, value:int=0, owner:HCQDeviceType|None=None, is_timeline:bool=False, timestamp_divider=1000):
+    self.base_buf, self.value_addr, self.timestamp_addr, self.owner = base_buf, base_buf.va_addr+0, base_buf.va_addr+8, owner
+    self.is_timeline = is_timeline
     self.timestamp_divider:decimal.Decimal = decimal.Decimal(timestamp_divider)
-    if isinstance(self.base_buf.va_addr, int) and not virt: self.value = value
+
+    if isinstance(self.base_buf.va_addr, int):
+      self.value_mv, self.timestamp_mv = self.base_buf.cpu_view().view(0, 8, 'Q'), self.base_buf.cpu_view().view(8, 8, 'Q')
+      self.value_mv[0] = value
 
   def __del__(self):
-    if self.should_return: HCQCompiled.signal_pool[unwrap(self.owner).peer_group].append(self.base_buf)
+    if isinstance(self.base_buf.va_addr, int) and self.owner is not None: HCQCompiled.signal_pool[self.owner.peer_group].append(self.base_buf)
 
   @property
-  def value_addr(self) -> sint: return self.base_buf.va_addr
-
-  @property
-  def timestamp_addr(self) -> sint: return self.base_buf.va_addr + 8
-
-  @property
-  def value(self) -> int: return self.base_buf.cpu_view().view(0, 8, 'Q')[0]
+  def value(self) -> int: return self.value_mv[0]
 
   @value.setter
-  def value(self, new_value:int): self.base_buf.cpu_view().view(0, 8, 'Q')[0] = new_value
+  def value(self, new_value:int): self.value_mv[0] = new_value
 
   @property
   def timestamp(self) -> decimal.Decimal:
@@ -245,7 +242,7 @@ class HCQSignal(Generic[HCQDeviceType]):
     Returns:
       The timestamp in microseconds.
     """
-    return self.base_buf.cpu_view().view(8, 8, 'Q')[0] / self.timestamp_divider
+    return self.timestamp_mv[0] / self.timestamp_divider
 
   def _sleep(self, time_spent_since_last_sleep_ms:int):
     """
@@ -304,8 +301,8 @@ class CLikeArgsState(HCQArgsState[ProgramType]):
 class HCQProgram(Generic[HCQDeviceType]):
   def __init__(self, args_state_t:Type[HCQArgsState], dev:HCQDeviceType, name:str, kernargs_alloc_size:int, lib:bytes|None=None, base:int|None=None):
     self.args_state_t, self.dev, self.name, self.kernargs_alloc_size = args_state_t, dev, name, kernargs_alloc_size
-    self.prof_prg_counter = next(self.dev.prof_prg_counter)
-    if PROFILE: Compiled.profile_events += [ProfileProgramEvent(dev.device, name, lib, base, self.prof_prg_counter)]
+    self.dev.prof_prg_counter += 1
+    if PROFILE: Compiled.profile_events += [ProfileProgramEvent(dev.device, name, lib, base, self.dev.prof_prg_counter)]
 
   @staticmethod
   def _fini(dev, buf, spec): dev.allocator.free(buf, buf.size, spec)
@@ -381,7 +378,7 @@ class HCQCompiled(Compiled, Generic[SignalType]):
     self.timeline_signal, self._shadow_timeline_signal = self.new_signal(value=0, is_timeline=True), self.new_signal(value=0, is_timeline=True)
     self.sig_prof_records:list[tuple[HCQSignal, HCQSignal, str|TracingKey, str]] = []
     self.prof_exec_counter:int = 0
-    self.prof_prg_counter = itertools.count(0)
+    self.prof_prg_counter:int = 0
 
     self.kernargs_buf:HCQBuffer = self.allocator.alloc(kernargs_size, BufferSpec(cpu_access=True))
     self.kernargs_offset_allocator:BumpAllocator = BumpAllocator(self.kernargs_buf.size, wrap=True)

tinygrad/runtime/support/system.py

@@ -346,7 +346,6 @@ class APLRemoteIfaceBase(LNXPCIIfaceBase):
       cls.gpus = System.pci_scan_bus(vendor, devices, base_class)
       if not cls.gpus: raise RuntimeError("No supported GPUs found")
       if not os.path.exists(APLRemotePCIDevice.APP_PATH): APLRemotePCIDevice.install_tinygpu()
-    if dev_id >= len(cls.gpus): raise RuntimeError(f"No device found for {dev_id}. Requesting more devices than the system has ({cls.gpus})?")
     self.pci_dev = APLRemotePCIDevice(dev.__class__.__name__[:2], f'remote:{dev_id}', bars)
     self.dev, self.vram_bar = dev, vram_bar
 

tinygrad/runtime/support/usb.py

@@ -1,7 +1,7 @@
 import ctypes, struct, dataclasses, array, itertools
 from typing import Sequence
 from tinygrad.runtime.autogen import libusb
-from tinygrad.helpers import DEBUG, to_mv, round_up, OSX, getenv
+from tinygrad.helpers import DEBUG, to_mv, round_up, OSX
 from tinygrad.runtime.support.hcq import MMIOInterface
 
 class USB3:
@@ -323,5 +323,3 @@ class USBMMIOInterface(MMIOInterface):
 
     _, acc_sz = self._acc_size(len(data) * struct.calcsize(self.fmt))
     self.usb.pcie_mem_write(self.addr+off, [int.from_bytes(data[i:i+acc_sz], "little") for i in range(0, len(data), acc_sz)], acc_sz)
-
-if getenv("MOCKGPU"): from test.mockgpu.usb import MockUSB3 as USB3  # type: ignore  # noqa: F811

tinygrad/schedule/indexing.py

@@ -3,7 +3,7 @@ import functools, itertools
 from dataclasses import dataclass, field
 from tinygrad.dtype import dtypes, AddrSpace
 from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, graph_rewrite, sint, AxisType, profile_matches
-from tinygrad.uop.ops import consumer_map_from_toposort, gate_kernel_sink
+from tinygrad.uop.ops import consumer_map_from_toposort, gate_kernel_sink, pm_gate_kernel_sink
 from tinygrad.uop.symbolic import symbolic, pm_simplify_valid, pm_drop_and_clauses
 from tinygrad.helpers import argsort, all_same, cpu_profile, PCONTIG, colored
 
@@ -21,7 +21,7 @@ def realize_assign_src(ctx:dict[UOp, None], buf:UOp, x:UOp):
   # you don't usually have to do this for assign unless there's a WAR hazard like TestAssign.test_assign_double_diamond_reduce
   if buf.base in x.backward_slice_with_self: ctx[x] = None
 
-pm_generate_realize_map = PatternMatcher([
+pm_generate_realize_map = pm_gate_kernel_sink+PatternMatcher([
   # always realize SINK src
   (UPat(Ops.SINK, name="s"), lambda ctx,s: ctx.update((x.base, None) for x in s.src if x.base.op not in ALWAYS_CONTIGUOUS)),
   # always realize
@@ -72,7 +72,7 @@ def create_bufferize_and_index_based_on_ranges(ctx:IndexingContext, x:UOp):
         # None in the device assigns it a number later
         opts = BufferizeOpts(device=s.device, removable=removable) if len(ctx.range_map[s][1]) == len(realized_ranges) else \
                BufferizeOpts(device=s.device, addrspace=AddrSpace.LOCAL, removable=removable)
-        new_src = UOp(Ops.BUFFERIZE, s.dtype, src=(new_src,)+closed_ranges, arg=opts)
+        new_src = UOp(Ops.BUFFERIZE, s.dtype, src=(new_src,)+closed_ranges, arg=opts, tag=s.tag if opts.addrspace == AddrSpace.GLOBAL else None)
         if x in ctx.range_map: new_src = new_src.index(*[r for i,r in enumerate(ctx.range_map[x][0]) if i in realized_ranges])
     new_srcs.append(new_src)
   # NOTE: do we need this?
@@ -88,7 +88,7 @@ def convert_pad_to_where_to_keep_behavior_local(ctx:IndexingContext, x:UOp):
 def convert_reduce_axis_to_reduce_with_ranges(ctx:IndexingContext, x:UOp):
   # input ranges
   new_ranges = [r for i,r in enumerate(ctx.range_map[x][0]) if i in x.arg[1]]
-  ret = UOp(Ops.REDUCE, x.dtype, src=(x.src[0],)+tuple(new_ranges), arg=x.arg[0])
+  ret = UOp(Ops.REDUCE, x.dtype, src=(x.src[0],)+tuple(new_ranges), arg=x.arg[0], tag=x.tag)
   ctx.range_map[ret] = ctx.range_map[x]
   return ret
 

tinygrad/schedule/multi.py

@@ -1,6 +1,6 @@
 import functools, itertools
-from tinygrad.helpers import all_same, all_int, prod, DEBUG, RING, ALL2ALL, getenv
-from tinygrad.uop.ops import Ops, UOp, PatternMatcher, UPat, GroupOp
+from tinygrad.helpers import all_same, all_int, prod, DEBUG, RING, ALL2ALL, VIZ, getenv
+from tinygrad.uop.ops import Ops, UOp, PatternMatcher, UPat, GroupOp, graph_rewrite_map, graph_rewrite
 from tinygrad.dtype import dtypes
 
 # *** allreduce implementation ***
@@ -187,3 +187,9 @@ multi_pm = PatternMatcher([
   (UPat(Ops.AFTER, src=(UPat(Ops.MULTI, name="multi"), UPat(Ops.CALL)), name="a"),
     lambda multi,a: a.replace(src=(multi.src[0],)+a.src[1:]).multi(multi.axis)),
 ])+replace_allreduce
+
+def get_multi_map(big_sink:UOp) -> dict[UOp, UOp]:
+  if VIZ: graph_rewrite(big_sink, PatternMatcher([]), name="View Multi AST")
+  ret = graph_rewrite_map(big_sink, multi_pm, name="multi_pm")
+  if VIZ: graph_rewrite(ret[big_sink], PatternMatcher([]), name="View Post Multi AST")
+  return ret

tinygrad/schedule/rangeify.py

@@ -1,15 +1,14 @@
 from dataclasses import dataclass, field, replace
 import itertools
 from tinygrad.dtype import dtypes, PtrDType, ImageDType, AddrSpace
-from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, _substitute, KernelInfo
-from tinygrad.uop.ops import graph_rewrite, sint, AxisType, BottomUpGate, profile_matches
+from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, _substitute, KernelInfo, pm_gate_kernel_sink
+from tinygrad.uop.ops import graph_rewrite, identity_element, sint, AxisType, BottomUpGate, _remove_all_tags
 from tinygrad.uop.symbolic import symbolic
-from tinygrad.helpers import prod, all_same, getenv, dedup, all_int, DEBUG, SPLIT_REDUCEOP, DEBUG_RANGEIFY, VIZ, MAX_KERNEL_BUFFERS
+from tinygrad.helpers import argsort, prod, all_same, getenv, flatten, dedup, all_int, DEBUG, SPLIT_REDUCEOP, DEBUG_RANGEIFY, VIZ, MAX_KERNEL_BUFFERS
 from tinygrad.helpers import PCONTIG, partition, get_single_element
 from tinygrad.codegen.simplify import pm_flatten_range, pm_reduce_simplify
 from tinygrad.codegen.opt import Opt
 from tinygrad.schedule.indexing import run_rangeify, BufferizeOpts, ALWAYS_CONTIGUOUS, IndexingContext, apply_movement_op
-from tinygrad.schedule.multi import multi_pm
 
 # creation can recurse a lot
 import sys
@@ -27,18 +26,34 @@ pm_mops = PatternMatcher([
    lambda r,idx: r.src[0].index(*apply_movement_op(r.op, r.src[0].shape, r.marg, idx.src[1:]), dtype=idx.dtype, arg=idx.arg)),
   # move movement ops after AFTER
   (UPat(GroupOp.Movement, name="r").after(name="a", allow_any_len=True),
-   lambda r,a: UOp(r.op, r.dtype, (a.replace(src=(r.src[0],)+a.src[1:]),)+r.src[1:], r.arg)),
+   lambda r,a: UOp(r.op, r.dtype, (a.replace(src=(r.src[0],)+a.src[1:], tag=None),)+r.src[1:], r.arg, tag=a.tag)),
   (UPat(GroupOp.Movement, name="r").end(name="a", allow_any_len=True), lambda r,a: a.replace(src=(r.src[0],)+a.src[1:])),
 ])
 
 # *****************
 # 0. do some cleanup rewrites, mostly copied from the old stuff
 
+def assign_to_contiguous(assign:UOp, target:UOp, src:UOp):
+  if (t := target.base).op is Ops.PARAM or (t.op is Ops.MSTACK and all(s.op is Ops.PARAM for s in t.src)): return None
+  # partial view of unrealized graph: insert CONTIGUOUS at base to realize it
+  if target is not t and target.op_in_backward_slice_with_self(Ops.SHRINK):
+    if t.op is Ops.CONTIGUOUS: return None
+    mops: list[UOp] = []
+    while target.op in GroupOp.Movement:
+      mops.append(target)
+      target = target.src[0]
+    new_target = t.f(Ops.CONTIGUOUS, tag=t.tag)
+    for m in reversed(mops): new_target = m.replace(src=(new_target,)+m.src[1:])
+    return assign.replace(src=(new_target, src))
+  return src.f(Ops.CONTIGUOUS, tag=assign.tag)
+
 def fix_assign_hazard(assign:UOp, target:UOp, src:UOp):
   # PERMUTE and FLIP reorder indices, SHRINK can have overlapping regions when dest is also shrunk
   unsafe = {Ops.PERMUTE, Ops.FLIP} | ({Ops.SHRINK} if target.op_in_backward_slice_with_self(Ops.SHRINK) else set())
-  if any(s.op in unsafe and target.base in s.backward_slice for s in src.toposort(gate=lambda s:s.op not in ALWAYS_CONTIGUOUS)):
-    return assign.replace(src=(target, src.contiguous()))
+  if not (hazards:=[s for s in src.toposort(gate=lambda s:s.op not in ALWAYS_CONTIGUOUS) if s.op in unsafe]): return
+  for h in hazards:
+    if any(s is target.base for s in h.toposort(gate=lambda s:s.op not in ALWAYS_CONTIGUOUS-{Ops.PARAM})):
+      return assign.replace(src=(target, src.contiguous()))
 
 def normalize_assign_target_chain(assign:UOp, target:UOp, src:UOp):
   root_target = target
@@ -68,54 +83,69 @@ def split_reduceop(reduce:UOp, x:UOp):
   splitted = x.reshape(splitted_shape).permute(tuple([d for d in range(len(splitted_shape)) if d!=dim_to_split]+[dim_to_split]))
   if DEBUG >= 3: print(f"split {divisor}: {x.shape} -> {splitted.shape} -> {reduce.shape}")
   # reduce original axes, then split
-  return splitted.r(*reduce.arg).contiguous().r(reduce.arg[0], (len(reduce.shape),)).reshape(reduce.shape)
+  return splitted.r(*reduce.arg).contiguous().r(reduce.arg[0], (len(reduce.shape),)).reshape(reduce.shape).replace(tag=reduce.tag)
 
 mop_cleanup = PatternMatcher([
-  # merge adjacent RESHAPES
-  (UPat(Ops.RESHAPE, src=(UPat(Ops.RESHAPE, name="x2"), UPat()), name="x"), lambda x,x2: x.replace(src=(x2.src[0], x.src[1]))),
+  # merge adjacent RESHAPES, safe because they are not tagged
+  (UPat(Ops.RESHAPE, src=(UPat(Ops.RESHAPE, name="x2"), UPat()), name="x"),
+   lambda x,x2: x.replace(src=(x2.src[0], x.src[1])) if x.tag is None and x2.tag is None else None),
 ])
 
-pm_gather_params = PatternMatcher([ (UPat(Ops.PARAM, name="p"), lambda ctx, p: ctx.append(p)), ])
-def resolve_call(c:UOp, allow_param_mismatch=False) -> UOp|None:
+def resolve_call(c:UOp) -> UOp|None:
   # don't resolve real kernel calls, sink or program
   if c.src[0].op is Ops.SINK and isinstance(c.src[0].arg, KernelInfo): return None
   if c.src[0].op is Ops.PROGRAM: return None
-  params: list[UOp] = []
-  graph_rewrite(c.src[0], pm_gather_params, bottom_up=True, ctx=params)
-  params = sorted(params, key=lambda x: x.arg)
+  params = sorted([x for x in c.src[0].toposort() if x.op == Ops.PARAM], key=lambda x: x.arg)
   args = c.src[1:]
   # TODO: this check belongs in spec, not here
-  if not allow_param_mismatch:
-    if [x.arg for x in params] != list(range(len(params))): raise RuntimeError(f"params not in order: {[x.arg for x in params]}")
-    if len(params) != len(args): raise TypeError(f"expected {len(params)} args, got {len(args)}")
+  if [x.arg for x in params] != list(range(len(params))): raise RuntimeError(f"params not in order: {[x.arg for x in params]}")
+  if len(params) != len(args): raise TypeError(f"expected {len(params)} args, got {len(args)}")
   for i, (p, a) in enumerate(zip(params, args)):
     if p.shape != a.shape: raise TypeError(f"arg {i} shape mismatch: expected {p.shape}, got {a.shape}")
     if p.dtype != a.dtype: raise TypeError(f"arg {i} dtype mismatch: expected {p.dtype}, got {a.dtype}")
-  return c.src[0].substitute(dict(zip(params, args)))
+  return c.src[0].substitute(dict(zip(params, args))).rtag(c.tag)
 
 earliest_rewrites = mop_cleanup+PatternMatcher([
+  # just removing it works...
+  (UPat((Ops.DETACH, Ops.CONTIGUOUS_BACKWARD), name="x"), lambda x: x.src[0]),
+
   # resolve calls
   (UPat(Ops.CALL, name="c"), resolve_call),
 
+  # remove CONTIGUOUS if the source is already contiguous
+  (UPat(Ops.RESHAPE, src=(UPat((Ops.PARAM, Ops.CONTIGUOUS)), UPat()), name="r").f(Ops.CONTIGUOUS, name="c"), lambda r,c: r.replace(tag=c.tag)),
+
   # split_reduceop
   (UPat(Ops.REDUCE_AXIS, name="reduce", src=(UPat.var("x"),)), split_reduceop),
 
+  # preserve tags?
+  # 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),
+
   # remove contiguous on movement ops before a copy on disk
   (UPat(GroupOp.Movement-{Ops.SHRINK, Ops.RESHAPE}, name="x").f(Ops.CONTIGUOUS).f(Ops.COPY, allow_any_len=True, name="copy"),
    lambda x,copy: copy.replace(src=(x,)+copy.src[1:]) if isinstance(x.device, str) and x.device.startswith("DISK") else None),
   # push copy past movement ops to disk
   (UPat(GroupOp.Movement-{Ops.SHRINK, Ops.RESHAPE}, name="x").f(Ops.COPY, allow_any_len=True, name="copy"),
-   lambda x,copy: x.replace(src=(copy.replace(src=(x.src[0],)+copy.src[1:]),)+x.src[1:]) \
+   lambda x,copy: x.replace(src=(copy.replace(src=(x.src[0],)+copy.src[1:], tag=None),)+x.src[1:], tag=copy.tag) \
       if isinstance(x.device, str) and x.device.startswith("DISK") else None),
 
   # ** copy rules **
 
+  # early fixup const copy
+  (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),
+
   # COPY and source size need to match
+  # TODO: expand after copy creates issues with tagging
   (UPat(Ops.COPY, src=(UPat(GroupOp.Movement, name="r"), UPat(name="d")), name="c"),
    lambda c,r,d: c.replace(src=(r.contiguous(), d)) if r.size != r.base.size else None),
 
   # copy only to different device
-  (UPat(Ops.COPY, src=(UPat.var("x"), UPat()), name="copy"), lambda x,copy: x.f(Ops.NOOP) if x.device == copy.device else None),
+  (UPat(Ops.COPY, src=(UPat.var("x"), UPat()), name="copy"), lambda x,copy: x.f(Ops.NOOP, tag=copy.tag) if x.device == copy.device else None),
 
   # ** assign rules **
 
@@ -123,20 +153,23 @@ earliest_rewrites = mop_cleanup+PatternMatcher([
   (UPat(Ops.ASSIGN, src=(UPat(name="target"), UPat(Ops.ASSIGN, src=(UPat(name="target"), UPat()), name="src"))), lambda target, src: src),
 
   # move bitcast from assign target to source: a.bitcast(X).assign(src) -> a.assign(src.bitcast(a.dtype))
-  (UPat(Ops.ASSIGN, src=(UPat(Ops.BITCAST, src=(UPat(name="target"),)), UPat(name="src"))),
-   lambda target, src: target.assign(src.bitcast(target.dtype))),
+  (UPat(Ops.ASSIGN, src=(UPat(Ops.BITCAST, src=(UPat(name="target"),)), UPat(name="src")), name="assign"),
+   lambda assign, target, src: target.assign(src.bitcast(target.dtype)).replace(tag=assign.tag)),
 
   # if assign target is itself an ASSIGN chain, canonicalize to the original buffer target
   (UPat(Ops.ASSIGN, src=(UPat(Ops.ASSIGN, name="target"), UPat(name="src")), allow_any_len=True, name="assign"), normalize_assign_target_chain),
 
-  # make source contiguous if it has hazardous movement ops on the dest buffer
-  (UPat(Ops.ASSIGN, src=(UPat.var("target"), UPat.var("src")), name="assign"), fix_assign_hazard),
+  # assign only to buffer, otherwise make it a CONTIGUOUS
+  (UPat(Ops.ASSIGN, src=(UPat(GroupOp.All-{Ops.PARAM}, name="target"), UPat(name="src")), name="assign"), assign_to_contiguous),
+
+   # make source contiguous if it has hazardous movement ops on the dest buffer
+   (UPat(Ops.ASSIGN, src=(UPat.var("target"), UPat.var("src")), name="assign"), fix_assign_hazard),
 ])
 
 # *****************
 # 3.5 cleanups
 
-ALWAYS_RUN_OPS = {Ops.CONTIGUOUS, Ops.COPY, Ops.ASSIGN, Ops.ENCDEC, Ops.NOOP}
+ALWAYS_RUN_OPS = {Ops.CONTIGUOUS, Ops.COPY, Ops.ASSIGN, Ops.ENCDEC}
 
 # you don't know in the first pass if axes are going to die, this happens if there's an EXPAND to the left
 def cleanup_dead_axes(b:UOp):
@@ -157,7 +190,8 @@ def cleanup_dead_axes(b:UOp):
       reshape.append(s)
       new_rng.append(rng)
   if hit:
-    return b.replace(src=b.src[0:1]+tuple(new_rng)).reshape(tuple(reshape)).expand(b.shape)
+    # move the tag to the expand. NOTE: this expand tag might not survive
+    return b.replace(src=b.src[0:1]+tuple(new_rng), tag=None).reshape(tuple(reshape)).expand(b.shape).replace(tag=b.tag)
 
 def gate_substitute(ctx, b:UOp) -> None:
   if not any(r in b.ranges for r in ctx.keys()): raise BottomUpGate()
@@ -230,7 +264,8 @@ def remove_bufferize(src:UOp, buf:UOp, idx:UOp):
 
 def remove_noop_bufferize(idx,b2):
   if idx.src[1:] != b2.src[1:] or idx.src[0].op is Ops.BUFFER_VIEW: return None
-  return idx.src[0].shrink(tuple((0, s) for s in b2.shape)) if b2.shape else idx.src[0]
+  new_tag = (idx.src[0].tag or ()) + (b2.tag or ()) or None
+  return idx.src[0].rtag(new_tag).shrink(tuple((0, s) for s in b2.shape)) if b2.shape else idx.src[0].rtag(new_tag)
 
 pm_const_buffer_folding = pm_mops+PatternMatcher([
   (UPat(Ops.BUFFERIZE, name="b"), cleanup_dead_axes),
@@ -240,13 +275,13 @@ pm_const_buffer_folding = pm_mops+PatternMatcher([
   # remove noop buffers. if we look at the next index we can remove even more of these
   (UPat(Ops.INDEX, name="idx").f(Ops.BUFFERIZE, allow_any_len=True, name="b2"), remove_noop_bufferize),
   # no buffers for const (ranges don't matter for const - it's the same value everywhere)
-  (UPat(Ops.CONST, name='c').f(Ops.BUFFERIZE, allow_any_len=True, name="b"), lambda c,b: b.const_like(c.arg)),
+  (UPat(Ops.CONST, name='c').f(Ops.BUFFERIZE, allow_any_len=True, name="b"), lambda c,b: b.const_like(c.arg).rtag(b.tag)),
   # indexing a const is a const
   (UPat(Ops.INDEX, src=(UPat(Ops.CONST, name="c"),),), lambda c: c),
   # copy on CONST is CONST
   (UPat(Ops.COPY, src=(UPat.cvar("x"), UPat()), name="copy"), lambda copy,x: copy.const_like(x.arg)),
   # hack if a noop turned to a const
-  (UPat(Ops.NOOP, src=(UPat.cvar("c"),), name="noop"), lambda c,noop: c),
+  (UPat(Ops.NOOP, src=(UPat.cvar("c"),), name="noop"), lambda c,noop: c.rtag(noop.tag)),
   # mstack on CONST is CONST
   (UPat(Ops.MSTACK, src=(UPat.var("s"),), allow_any_len=True).f(Ops.INDEX, allow_any_len=True),
    lambda s: UOp.const(c.dtype, c.arg) if (c:=s.base).op is Ops.CONST else None),
@@ -272,7 +307,7 @@ def late_buffer_view(t:UOp, b:UOp):
   if len(shape) == 0: offset = x.src[1].arg
   else: offset = max(sum(idx.vmin for idx in x.src[1:]), 0)
 
-  return b.replace(src=(UOp(Ops.BUFFER_VIEW, t.dtype, (x.base,), (size, offset)), b.src[1]))
+  return b.replace(src=(UOp(Ops.BUFFER_VIEW, t.dtype, (x.base,), (size, offset), tag=t.tag), b.src[1]))
 
 to_bufferview = PatternMatcher([
   (UPat(Ops.BUFFERIZE, src=(UPat((Ops.BITCAST, Ops.CONTIGUOUS), name="t"), UPat()), name="b"), late_buffer_view),
@@ -311,6 +346,7 @@ pm_limit_bufs = PatternMatcher([(UPat(set.union(GroupOp.Binary, GroupOp.Ternary)
 # NOTE: this has been fixed up a bit
 
 def bufferize_to_store(ctx:itertools.count, x:UOp, idx:UOp, allow_locals=True):
+  #assert isinstance(x.tag, Flat), "bufferize must be flat"
   size = prod(x.shape)
   rngs = sorted(idx.ranges, key=lambda x: x.arg)
   assert size > 0 and isinstance(size, int), f"no zero sized or symbolic sized buffers {size}"
@@ -323,14 +359,14 @@ def bufferize_to_store(ctx:itertools.count, x:UOp, idx:UOp, allow_locals=True):
     # skip self-assign from same-device copy, otherwise create the store
     # in assign, this is the buffer size, not the bufferize size
     if assign_src is assign_target: ret = assign_target.src[0]
-    else: ret = assign_target.src[0].after(assign_target.replace(dtype=sdtype).store(assign_src).end(*rngs))
+    else: ret = assign_target.src[0].after(assign_target.replace(dtype=sdtype).store(assign_src, tag=x.tag).end(*rngs))
     for op, marg in reversed(assign.arg or ()): ret = ret._mop(op, marg)
     return ret
 
   # NOTE: the DEFINE_LOCAL needs to be disambiguated here
   if sdtype.addrspace == AddrSpace.GLOBAL:
     buf = UOp(Ops.BUFFER, x.dtype, (UOp(Ops.LUNIQUE, arg=next(ctx)), UOp(Ops.DEVICE, arg=x.arg.device)), size)
-    do_store = buf.index(idx, dtype=sdtype).store(x.src[0]).end(*rngs)
+    do_store = buf.index(idx, dtype=sdtype).store(x.src[0], tag=x.tag).end(*rngs)
     return buf.after(do_store)
 
   if allow_locals:
@@ -339,16 +375,16 @@ def bufferize_to_store(ctx:itertools.count, x:UOp, idx:UOp, allow_locals=True):
     do_store = buf.broadcast(x.src[1].dtype.count).index(idx, dtype=sdtype).store(x.src[0]).end(*rngs)
     return buf.after(do_store.barrier())
 
-# collapse any BUFFERIZE to single input BUFFERIZE
+# collapse any BUFFERIZE to single input BUFFERIZE. move the tag to a reshape
 def flatten_bufferize(x:UOp):
-  if len(x.src) == 2: return None
-  ret = x.replace(src=(x.src[0], get_single_element(apply_movement_op(Ops.RESHAPE, (prod(x.shape),), x.shape, x.src[1:]))))
+  if x.tag is None and len(x.src) == 2: return None
+  ret = x.replace(tag=None, src=(x.src[0], get_single_element(apply_movement_op(Ops.RESHAPE, (prod(x.shape),), x.shape, x.src[1:]))))
   rngs = x.src[1:]
-  ret = ret.reshape(x.shape)
+  ret = ret.forced_reshape(x.shape)
   if any(r.op is Ops.RANGE and r.src[0].op is not Ops.CONST for r in rngs):
     sym_shape = tuple([r.src[0] if r.op is not Ops.CONST else 1 for r in rngs])
     ret = ret.shrink(tuple([(0,x) for x in sym_shape]))
-  return ret
+  return ret.rtag(x.tag)
 pm_flatten_bufferize = PatternMatcher([(UPat(Ops.BUFFERIZE, name="x"), flatten_bufferize)])
 
 pm_add_buffers = pm_mops+pm_flatten_bufferize+to_bufferview+PatternMatcher([
@@ -356,7 +392,7 @@ pm_add_buffers = pm_mops+pm_flatten_bufferize+to_bufferview+PatternMatcher([
 
   # move RESHAPEs through MSELECT/MSTACK
   (UPat((Ops.MSELECT, Ops.MSTACK), src=UPat(Ops.RESHAPE), name="m"),
-   lambda m: m.replace(src=tuple([x.src[0].base for x in m.src])).reshape(m.shape)),
+   lambda m: m.replace(src=tuple([x.src[0].base for x in m.src]), tag=None).reshape(m.shape).rtag(m.tag)),
 
   # remove any RESHAPEs on KERNEL
   (UPat(Ops.CALL, name="k"), lambda k: k.replace(src=tuple(x.src[0] if x.op is Ops.RESHAPE else x for x in k.src))),
@@ -375,6 +411,7 @@ class LocalAddBufferContext:
   map:dict = field(default_factory=dict)
   vars:dict = field(default_factory=dict)
   range:int = 0
+  parent_tags:list = field(default_factory=list)
   opts:tuple|None = None
 
 def debuf(ctx:LocalAddBufferContext, buf:UOp):
@@ -412,10 +449,6 @@ to_define_global = PatternMatcher([
   (UPat(Ops.STORE, name="x"), find_bufs),
   (UPat(Ops.BUFFER, name="buf"), debuf),
   (UPat(Ops.PARAM, src=(UPat(), UPat(Ops.DEVICE)), name="buf"), debuf),
-  (UPat(Ops.PARAM, src=(UPat(), UPat(), UPat.cvar('vmin'), UPat.cvar('vmax'), UPat.var("nm")), name="v"),
-   lambda v, vmin, vmax, nm: UOp.variable(nm.arg, vmin.arg, vmax.arg, v.dtype)),
-  (UPat(Ops.INDEX, src=(UPat(Ops.DEFINE_VAR, name="v"),)), lambda v: v),
-
   (UPat(Ops.BIND, name="b"), unbind_kernel),
   (UPat((Ops.MSTACK, Ops.MSELECT, Ops.AFTER), name="after"), handle_after),
 
@@ -438,7 +471,13 @@ rangeify_codegen = PatternMatcher([
 
   # no NOOP in the kernel graph
   # TODO: this can be moved into codegen?
-  (UPat(Ops.NOOP, name="x"), lambda x: x.src[0] if len(x.src) else None),
+  (UPat(Ops.NOOP, name="x"), lambda x: x.src[0]),
+
+  # add loads to non ptr indexes
+  # TODO: this can be moved into codegen?
+  #(UPat.any(UPat(Ops.DEFINE_GLOBAL, name="dg"), UPat(Ops.DEFINE_LOCAL).f(Ops.AFTER, allow_any_len=True, name="dg"))
+  # .f(Ops.INDEX, name="idx", allow_any_len=True),
+  #  lambda dg,idx: None if isinstance(idx.dtype, (PtrDType, ImageDType)) else idx.replace(dtype=dg.dtype, arg=None).load()),
 
   # fix broadcast dtype
   (UPat(Ops.AFTER, name="a").broadcast(name="b"), lambda a,b: a.broadcast(len(b.src))),
@@ -451,17 +490,29 @@ rangeify_codegen = PatternMatcher([
       idx.replace(dtype=dg.dtype, arg=None).load(dtype=dg.dtype.base.scalar().vec(dg.dtype.vcount))),
 ])
 
+def remove_metadata_tags(ctx:LocalAddBufferContext, x:UOp):
+  if x.tag is None or x.tag == (): return None
+  if isinstance(x.tag, tuple): ctx.parent_tags += list(x.tag)
+  return x.replace(tag=None)
+
+pm_remove_tags = PatternMatcher([
+  (UPat(GroupOp.All, name="x"), remove_metadata_tags),
+])
+
 pm_add_range_tags = PatternMatcher([
   (UPat(Ops.RANGE, name="x"), lambda x: x.rtag(())),
 ])
 
-def split_store(x:UOp) -> UOp|None:
+def split_store(ctx:list[UOp], x:UOp) -> UOp|None:
   # if we have any open ranges here, we don't split
   if x.ranges: return None
 
   # local kernel rewrite
   lctx = LocalAddBufferContext()
-  ret = graph_rewrite(x, to_define_global+pm_flatten_range+rangeify_codegen, ctx=lctx, name="kernel split", bottom_up=True)
+  ret = graph_rewrite(x, to_define_global+pm_flatten_range+rangeify_codegen+pm_remove_tags, ctx=lctx, name="kernel split", bottom_up=True)
+
+  # gather the metadata
+  metadatas = [ctx[y].metadata for y in lctx.parent_tags]
 
   # SINK requires all buffers on the same device, but COPY/BUFFER_VIEW/ENCDEC are cross-device or special hardware ops
   if ret.op is Ops.STORE: stored = ret.src[1]
@@ -470,7 +521,8 @@ def split_store(x:UOp) -> UOp|None:
   if stored.op in {Ops.COPY, Ops.BUFFER_VIEW, Ops.ENCDEC}: ret = stored
   else: ret = ret.sink(arg=KernelInfo(opts_to_apply=lctx.opts))
 
-  kernel = ret.call(*lctx.map.values(), *lctx.vars.keys())
+  metadata = tuple(dedup(flatten([x for x in metadatas if x is not None])))[::-1]
+  kernel = ret.call(*lctx.map.values(), *lctx.vars.keys(), metadata=metadata)
   if ret.op is Ops.SINK and not all_same([x.device for x in kernel.src[1:] if x.op is not Ops.BIND]):
     raise RuntimeError(f"all buffers must be on the same device: {tuple(b.buf_uop for b in kernel.src[1:])}")
   return kernel
@@ -479,10 +531,42 @@ split_kernels = PatternMatcher([
   (UPat((Ops.STORE, Ops.END), name="x"), split_store),
 ])
 
-@profile_matches
-def get_kernel_graph(sink:UOp) -> UOp:
-  tsink = graph_rewrite(sink, multi_pm, name="multi_pm", rewrite_into_calls=True)
-  tsink = graph_rewrite(tsink, pm_syntactic_sugar+pm_mops+earliest_rewrites, bottom_up=True, name="earliest rewrites")
+def tag_uop(ctx:tuple[list[UOp], set[UOp]], x:UOp):
+  if x.tag is not None or x in ctx[1]: return None
+  if x.tag is None and x.op is Ops.CALL:
+    # don't tag anything in a CALL
+    for u in x.src[0].toposort(): ctx[1].add(u)
+  if x.dtype.scalar() == dtypes.index: return None
+  ctx[0].append(x)
+  return x.replace(tag=(len(ctx[0])-1,))
+add_tags = pm_gate_kernel_sink+PatternMatcher([
+  # don't tag BUFFERs, they are global
+  (UPat(GroupOp.All-{Ops.PARAM, Ops.CONST, Ops.DEVICE, Ops.UNIQUE, Ops.LUNIQUE, Ops.DEFINE_VAR, Ops.BIND, Ops.END,
+                     Ops.MSTACK, Ops.MSELECT, Ops.RANGE}.union(GroupOp.Movement), name="x"), tag_uop),
+  (UPat({Ops.MSTACK, Ops.MSELECT}, name="x"), lambda ctx,x: None if all(s.op is Ops.PARAM for s in x.src) else tag_uop(ctx, x)),
+])
+
+# support for using a contiguous permuted view instead of the parent view if one exists
+
+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
+replace_contiguous = PatternMatcher([
+  (UPat(Ops.CONTIGUOUS, src=(UPat(GroupOp.Movement, name="src"),), name="contig"), found_contiguous),
+  (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),
+])
+
+def get_rangeify_map(sink:UOp) -> dict[UOp, UOp]:
+  if VIZ: graph_rewrite(sink, PatternMatcher([]), name="View Input Graph")
+  uop_list: list[UOp] = []
+  tsink = graph_rewrite(sink, add_tags, ctx=(uop_list, set()), bottom_up=True, name="number the uops")
+
+  tsink = graph_rewrite(tsink, pm_syntactic_sugar+pm_mops+earliest_rewrites+replace_contiguous, ctx={}, bottom_up=True, name="earliest rewrites")
 
   # convert movement ops to ranges
   tsink, rctx = run_rangeify(tsink, bool(DEBUG_RANGEIFY))
@@ -490,12 +574,19 @@ def get_kernel_graph(sink:UOp) -> UOp:
   tsink = graph_rewrite(tsink, symbolic+pm_reduce_simplify+pm_const_buffer_folding+pm_remove_bufferize, name="symbolic+reduce_collapse+debuf")
   tsink = graph_rewrite(tsink, pm_limit_bufs, ctx=rctx, name="limit buffers")
 
-  if VIZ: graph_rewrite(tsink, PatternMatcher([]), name="View Rangeify")
+  # rebuild the sink with all the BUFFERIZEs with tags, this is what's ending up in the tensor graph
+  # MSTACK stacks multiple BUFFERIZEs in one tagged tensor
+  # if it's not tagged by here, it's out
+  tsink = UOp.sink(*[x for x in tsink.backward_slice if x.base.op in {Ops.BUFFERIZE, Ops.MSTACK, Ops.CONST, Ops.PARAM, Ops.AFTER} and \
+                     x.tag is not None and len(x.tag)])
+
+  if VIZ: graph_rewrite(tsink, PatternMatcher([]), name="View Tagged Rangeify")
 
   # bufferize -> store
   lunique_start: int = max([-1]+[x.arg for x in tsink.toposort() if x.op is Ops.LUNIQUE]) + 1
-  tsink = graph_rewrite(tsink, pm_add_buffers+pm_add_range_tags, ctx=itertools.count(lunique_start), bottom_up=True, name="bufferize to store")
-  tsink = graph_rewrite(tsink, split_kernels, bottom_up=True, name="split kernels")
+  tsink = graph_rewrite(tsink, pm_gate_kernel_sink+pm_add_buffers+pm_add_range_tags, ctx=itertools.count(lunique_start), bottom_up=True,
+                        name="bufferize to store")
+  tsink = graph_rewrite(tsink, pm_gate_kernel_sink+split_kernels, ctx=uop_list, bottom_up=True, name="split kernels")
 
   # WAR deps: if kernel U reads buffer S, and S is also written by another kernel, S's write must wait for U to finish
   afters = [u for u in tsink.toposort() if u.op is Ops.AFTER]
@@ -506,9 +597,21 @@ def get_kernel_graph(sink:UOp) -> UOp:
       # TODO: this is probably broken for MSELECT/MSTACK
       if s.op not in {Ops.BUFFER, Ops.PARAM} or s is u.buf_uop or (a:=kernel_assign.get(s)) is None: continue
       if a.src[1] is u.src[1]: continue  # same kernel (multi-output custom kernels)
-      if any(x.op is Ops.AFTER and x.buf_uop is s for x in kernel_assign[u.buf_uop].backward_slice):
-        raise RuntimeError(f"cycle detected in assign graph, buffers {s} and {u.buf_uop} have circular dependency")
+      if any(x.op is Ops.AFTER and x.buf_uop is s for x in u.toposort()):
+        raise RuntimeError(f"cycle detected in graph, kernel for {u.buf_uop} must either depend on AFTER or BUFFER")
       assign_rep[a] = kernel_assign[s] = a.replace(src=a.src+(u,))
   if assign_rep: tsink = graph_rewrite(tsink, _substitute, ctx=assign_rep, bottom_up=True, name="fix_assign")
+
+  # TODO: we can probably get this earlier
+  sink_tags = [s.tag for s in tsink.src]
+  tsink = graph_rewrite(tsink, _remove_all_tags, name="remove all tags")
+
   if VIZ: graph_rewrite(tsink, PatternMatcher([]), name="View Kernel Graph")
-  return tsink
+
+  becomes_map: dict[UOp, UOp] = {}
+  for tag, s in zip(sink_tags, tsink.src):
+    assert tag is not None
+    for a in tag:
+      if a is None: continue
+      becomes_map[uop_list[int(a)]] = s
+  return becomes_map

tinygrad/tensor.py

@@ -13,11 +13,9 @@ from tinygrad.gradient import compute_gradient
 from tinygrad.mixin import OpMixin
 from tinygrad.mixin.movement import _align_left
 from tinygrad.uop.ops import smax, smin, resolve, UOp, Ops, sint, identity_element, all_metadata, _index_to_concrete_int, sint_to_uop, Variable
-from tinygrad.uop.ops import PatternMatcher, UPat
 from tinygrad.engine.schedule import ExecItem, complete_create_schedule_with_vars
 from tinygrad.device import Device, Buffer
 from tinygrad.engine.realize import run_schedule
-from tinygrad.engine.allocations import transform_to_call
 
 # TODO: this should be the only usage of Device
 def canonicalize_device(device:str|tuple|list|None) -> str|tuple[str, ...]:
@@ -27,8 +25,7 @@ def canonicalize_device(device:str|tuple|list|None) -> str|tuple[str, ...]:
 
 all_tensors: dict[weakref.ref[Tensor], None] = {}
 _pending_assigns: dict[UOp, list[UOp]] = {}  # buffer_uop -> [assign_uops in insertion order]
-_pm_strip_after_noop = PatternMatcher([(UPat(Ops.AFTER, src=(UPat(name="buf"), UPat(Ops.NOOP))), lambda ctx,buf: buf)])
-def _apply_map_to_tensors(applied_map:dict[UOp, UOp], name:str, extra_pm:PatternMatcher|None=None) -> None:
+def _apply_map_to_tensors(applied_map:dict[UOp, UOp], name:str) -> None:
   with cpu_profile(TracingKey(name), "TINY"):
     # get tensors in scope
     in_scope: dict[UOp, bool] = {}
@@ -37,7 +34,7 @@ def _apply_map_to_tensors(applied_map:dict[UOp, UOp], name:str, extra_pm:Pattern
 
     # get all Tensors and apply the map
     sink = UOp.sink(*[t.uop for t in scope_tensors])
-    new_sink = sink.substitute(applied_map, name=f"substitute {name}", extra_pm=extra_pm)
+    new_sink = sink.substitute(applied_map, name=f"substitute {name}")
 
     # set the relevant uop to the realized UOps
     for t,s,ns in zip(scope_tensors, sink.src, new_sink.src):
@@ -252,25 +249,17 @@ class Tensor(OpMixin):
     """
     return [Tensor(u, device=u.device) for u in UOp.custom_kernel(*[t.uop for t in (self,)+lst], fxn=fxn, grad_fxn=grad_fxn)]
 
-  def callify(self, *lst:Tensor) -> Tensor:
-    big_sink = UOp.sink(*[x.uop for x in (self,)+lst])
-    big_sink, buffer_map = transform_to_call(big_sink)
-    _apply_map_to_tensors({x:y.after(big_sink) for x,y in buffer_map.items()}, name="callify")
-    return self
-
   def schedule_with_vars(self, *lst:Tensor) -> tuple[list[ExecItem], dict[str, int]]:
     """
     Creates the schedule needed to realize these Tensor(s), with Variables.
 
     NOTE: A Tensor can only be scheduled once.
     """
-    # collect existing CALLs before callify (so we can clean them up in other tensors that share them)
-    pre_calls = {u for t in (self,)+lst for u in t.uop.toposort() if u.op is Ops.CALL}
-    self.callify(*lst)
-    calls, schedule, var_vals = complete_create_schedule_with_vars(UOp.sink(*[x.uop for x in (self,)+lst]))
-    # replace scheduled CALLs with NOOP so AFTER(buf, CALL) -> AFTER(buf, NOOP) -> buf in scope tensors
-    # include pre-existing CALLs too (they were reconstructed inside callify, but other tensors still reference the originals)
-    _apply_map_to_tensors({c:UOp(Ops.NOOP) for c in set(calls) | pre_calls}, name="buffers", extra_pm=_pm_strip_after_noop)
+    big_sink = UOp.sink(*[x.uop for x in (self,)+lst])
+
+    # this is where the schedule cache should go
+    becomes_map, schedule, var_vals = complete_create_schedule_with_vars(big_sink)
+    _apply_map_to_tensors(becomes_map, name="Apply Schedule Map")
     return schedule, var_vals
 
   def schedule(self, *lst:Tensor) -> list[ExecItem]:
@@ -289,12 +278,8 @@ class Tensor(OpMixin):
           # recursively realize pending assigns that this assign's value depends on
           for u in assign_uop.toposort():
             if u.op is Ops.BUFFER and u in _pending_assigns: _realize_pending(u)
-          sink = UOp.sink(assign_uop)
-          call, buffer_map = transform_to_call(sink)
-          callified_sink = UOp.sink(*[buffer_map.get(s, s).after(call) for s in sink.src])
-          calls, schedule, var_vals = complete_create_schedule_with_vars(callified_sink)
-          becomes_map = {**buffer_map, **{c:UOp(Ops.NOOP) for c in calls}}
-          _apply_map_to_tensors(becomes_map, name="Apply Pending Assign", extra_pm=_pm_strip_after_noop)
+          becomes_map, schedule, var_vals = complete_create_schedule_with_vars(UOp.sink(assign_uop))
+          _apply_map_to_tensors(becomes_map, name="Apply Pending Assign")
           run_schedule(schedule, var_vals, do_update_stats=do_update_stats)
           # update remaining pending assigns so they reference realized buffers instead of stale lazy graphs
           if becomes_map:
@@ -419,7 +404,7 @@ class Tensor(OpMixin):
     """
     Creates a clone of this tensor allocating a separate buffer for the data.
     """
-    ret = self.empty_like()
+    ret = Tensor.empty(self.shape, device=self.device, dtype=self.dtype)
     if self.grad is not None: ret.grad = self.grad.clone()
     return ret.assign(self)
 
@@ -552,15 +537,12 @@ class Tensor(OpMixin):
     device = canonicalize_device(device)
     return Tensor(UOp.new_buffer(device, size, dtype), device, dtype, **kwargs).shrink(((0,prod(shape)),)).reshape(shape)
 
-  def empty_like(self, dtype:DTypeLike|None=None, device:str|tuple[str, ...]|None=None, **kwargs) -> Tensor:
+  def empty_like(self, **kwargs) -> Tensor:
     """
     Creates an empty tensor with the same shape as `self`.
     If `dtype` is not specified, the dtype of `self` is used.
     """
-    dtype, device = self.dtype if dtype is None else dtype, self.device if device is None else device
-    if isinstance(device, tuple) and (axis := self.uop.axis) is not None:
-      return Tensor(Tensor.empty(self.uop.max_shard_shape, dtype=dtype, device=device, **kwargs).uop.multi(axis), device=device)
-    return Tensor.empty(self.shape, dtype=dtype, device=device, **kwargs)
+    return Tensor.empty(self.shape, dtype=kwargs.pop("dtype", self.dtype), device=kwargs.pop("device", self.device), **kwargs)
 
   @staticmethod
   def from_blob(ptr:int, shape:tuple[int, ...], **kwargs) -> Tensor:
@@ -646,7 +628,7 @@ class Tensor(OpMixin):
       Tensor._device_seeds[device] = Tensor(
         [int.from_bytes(hashlib.sha256(len(Tensor._device_seeds).to_bytes(4, "big")).digest(), "big"), Tensor._seed],
         device=device, dtype=dtypes.uint32, requires_grad=False)
-      Tensor._device_rng_counters[device] = Tensor([num], device=device, dtype=dtypes.uint32, requires_grad=False).contiguous()
+      Tensor._device_rng_counters[device] = Tensor([num], device=device, dtype=dtypes.uint32, requires_grad=False)
     # increment rng counter for devices
     else: Tensor._device_rng_counters[device].assign(Tensor._device_rng_counters[device] + num)
 
@@ -1093,34 +1075,6 @@ class Tensor(OpMixin):
 
   def _mop(self, op:Ops, arg) -> Tensor: return self._apply_uop(UOp._mop, extra_args=(op,), arg=arg)
 
-  def _pad_constant(self, pX:tuple[tuple[sint, sint], ...], value:float) -> Tensor:
-    # shrink first for negative pads, then pad with only non-negative values
-    has_neg = not all(resolve(p >= 0) for p in flatten(pX))
-    X = self.shrink(tuple((-smin(pB,0),smin(pA+s,s)) for (pB,pA),s in zip(pX, self.shape))) if has_neg else self
-    pads = tuple((smax(pB,0), smax(pA,0)) for pB,pA in pX) if has_neg else pX
-    if value == 0: return X._apply_uop(UOp.pad, arg=pads)
-    return X._apply_uop(UOp.pad, arg=pads) + Tensor.ones_like(X)._apply_uop(UOp.pad, arg=pads).where(0, value)
-
-  def _pad_circular(self, pX:tuple[tuple[sint, sint], ...]) -> Tensor:
-    if any(pB>sh or pA>sh for (pB,pA),sh in zip(pX, self.shape)): raise ValueError('Padding value causes wrapping around more than once.')
-    if any(pB<0 or pA<0 for pB,pA in pX): raise NotImplementedError("Negative pads with circular pads is not supported")
-    orig_shape, X = self.shape, self.repeat(tuple(1 + bool(pB) + bool(pA) for pB,pA in pX))
-    return X.shrink(tuple((0 if pB == 0 else osh-pB, xsh if pA == 0 else xsh-osh+pA) for (pB,pA),osh,xsh in zip(pX, orig_shape, X.shape)))
-
-  def _pad_reflect_replicate(self, pX:tuple[tuple[sint, sint], ...], mode:str) -> Tensor:
-    X, pads = self, tuple((smax(pB,0), smax(pA,0)) for pB,pA in pX)
-    for d,(pB,pA) in enumerate(pads):
-      if mode == "reflect":
-        if pB >= (s:=X.shape[d]) or pA>=s: raise ValueError(f"Padding ({pB}, {pA}) should be less than the input size={s} for dim={d}.")
-        slcB, slcA = slice(pB,0,-1), slice(s-2 if s-2>=0 else None, s-2-pA if s-2-pA>=0 else None, -1)
-        xB, xA = (X[[slc if i == d else slice(None) for i in range(X.ndim)]] if p > 0 else None for slc, p in ((slcB, pB), (slcA, pA)))
-      else:
-        shrB, shrA = tuple((0,1) if i==d else None for i in range(X.ndim)), tuple((X.shape[i]-1,X.shape[i]) if i==d else None for i in range(X.ndim))
-        xB, xA = (X.shrink(shr).expand(tuple(p if i==d else None for i in range(X.ndim))) if p > 0 else None for shr, p in ((shrB, pB), (shrA, pA)))
-      X = Tensor.cat(*(X_ for X_ in (xB, X, xA) if X_ is not None), dim=d)
-    # shrink after for negative pads (reflection/replication must see full data first)
-    return X.shrink(tuple((-min(pB,0), min(pA+s,s)) for (pB,pA),s in zip(pX, X.shape)))
-
   def pad(self, padding:Sequence[sint]|Sequence[tuple[sint, sint]|None], mode:str="constant", value:float=0.0) -> Tensor:
     """
     Returns a tensor with padding applied based on the input `padding`.
@@ -1153,18 +1107,36 @@ class Tensor(OpMixin):
     print(t.pad((1, 2, 0, -1), value=-float('inf')).numpy())
     ```
     """
-    # normalize to grouped format
+    if mode not in {"constant", "reflect", "replicate", "circular"}: raise NotImplementedError(f"{mode=} is not supported")
+    # flat padding
     if all(isinstance(p, (int,UOp)) for p in padding):
       if len(padding)%2 != 0: raise ValueError("Flat padding must have even number of pads")
       pX = _flat_to_grouped(tuple(cast(Sequence[sint], padding)) + (0,0)*(self.ndim - len(padding)//2))
+    # group padding
     else: pX = tuple((0,0) if p is None else p for p in cast(Sequence[tuple[sint, sint]|None], padding))
     if len(pX) != self.ndim: raise ValueError(f"padding length is improper, {padding=} {self.ndim=}")
-    # dispatch
-    if mode == "constant": return self._pad_constant(pX, value)
+    X, pads = self, tuple((smax(pB,0), smax(pA,0)) for pB,pA in pX)
+    if mode == "constant":
+      def _constant(x:Tensor,px,v) -> Tensor:
+        return x._apply_uop(UOp.pad, arg=px) if v == 0 else (x._apply_uop(UOp.pad, arg=px)+Tensor.ones_like(x)._apply_uop(UOp.pad, arg=px).where(0,v))
+      return _constant(X, pX, value) if all(resolve(p >= 0) for p in flatten(pX)) else \
+             _constant(X.shrink(tuple((-smin(pB,0),smin(pA+s,s)) for (pB,pA),s in zip(pX, X.shape))), pads, value)
     assert all_int(self.shape), f"does not support symbolic shape {self.shape}"
-    if mode == "circular": return self._pad_circular(pX)
-    if mode in {"reflect", "replicate"}: return self._pad_reflect_replicate(pX, mode)
-    raise NotImplementedError(f"{mode=} is not supported")
+    if mode == "circular":
+      if any(pB>sh or pA>sh for (pB,pA),sh in zip(pX, X.shape)): raise ValueError('Padding value causes wrapping around more than once.')
+      if any(pB<0 or pA<0 for pB,pA in pX): raise NotImplementedError("Negative pads with circular pads is not supported")
+      orig_shape, X = X.shape, X.repeat(tuple(1 + bool(pB) + bool(pA) for pB,pA in pads))
+      return X.shrink(tuple((0 if pB == 0 else osh-pB, xsh if pA == 0 else xsh-osh+pA) for (pB,pA),osh,xsh in zip(pads, orig_shape, X.shape)))
+    for d,(pB,pA) in enumerate(pads):
+      if mode == "reflect":
+        if pB >= (s:=X.shape[d]) or pA>=s: raise ValueError(f"Padding ({pB}, {pA}) should be less than the input size={s} for dim={d}.")
+        slcB, slcA, = slice(pB,0,-1), slice(s-2 if s-2>=0 else None, s-2-pA if s-2-pA>=0 else None, -1)
+        xB, xA = (X[[slc if i == d else slice(None) for i in range(X.ndim)]] if p > 0 else None for slc, p in ((slcB, pB), (slcA, pA)))
+      if mode == "replicate":
+        shrB, shrA, = tuple((0,1) if i==d else None for i in range(X.ndim)), tuple((X.shape[i]-1,X.shape[i]) if i==d else None for i in range(X.ndim))
+        xB, xA = (X.shrink(shr).expand(tuple(p if i==d else None for i in range(X.ndim))) if p > 0 else None for shr, p in ((shrB, pB), (shrA, pA)))
+      X = Tensor.cat(*(X_ for X_ in (xB, X, xA) if X_ is not None), dim=d)
+    return X.shrink(tuple((-min(pB,0), min(pA+s,s)) for (pB,pA),s in zip(pX, X.shape)))
 
   # convenience
   def pad_to(self, shape, *args):

tinygrad/uop/decompositions.py

@@ -335,6 +335,7 @@ def l2i(op: Ops, dt: DType, *uops:UOp):
     case Ops.CAST if dt in dtypes.floats:
       small = (a1.eq(0) & (a0 >= 0)) | (a1.eq(-1) & (a0 < 0))
       return small.where(a0.cast(dt), ((a1.cast(dtypes.float32) * (2**32)) + a0.bitcast(dtypes.uint).cast(dtypes.float32)).cast(dt))
+    case Ops.CAST if dt == dtypes.bool: return a0.ne(UOp.const(a0.dtype, 0)) | a1.ne(UOp.const(a1.dtype, 0))
     case Ops.CAST: return a0.bitcast(dtypes.uint).cast(dt)
     case Ops.BITCAST: return a0.bitcast(dt), a1.bitcast(dt)
     case Ops.SHL:

tinygrad/uop/ops.py

@@ -8,7 +8,7 @@ from tinygrad.dtype import ConstType, ImageDType, dtypes, DType, truncate, PtrDT
 from tinygrad.dtype import storage_fmt_for_dtype, to_storage_scalar, from_storage_scalar
 from tinygrad.helpers import ContextVar, all_int, prod, getenv, all_same, Context, partition, temp, unwrap, T, argfix, Metadata, flatten, TRACEMETA
 from tinygrad.helpers import PROFILE, dedup, cdiv, cmod, diskcache_put, to_function_name, cpu_profile, TracingKey, VIZ, SPEC, CAPTURE_PROCESS_REPLAY
-from tinygrad.helpers import strip_parens, colored, ansilen, printable
+from tinygrad.helpers import strip_parens, colored, ansilen, printable, panic
 if TYPE_CHECKING:
   from tinygrad.device import Buffer, MultiBuffer
   from tinygrad.renderer import Estimates
@@ -207,15 +207,10 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
     match self.op:
       # late ops don't have shape
       case Ops.UNIQUE | Ops.LUNIQUE | Ops.DEVICE | Ops.RANGE | Ops.LOAD | Ops.IF | Ops.BARRIER | Ops.CUSTOM | Ops.CUSTOMI | \
-           Ops.VECTORIZE | Ops.GEP | Ops.SPECIAL | Ops.UNROLL | Ops.CONTRACT | Ops.SINK | \
+           Ops.VECTORIZE | Ops.VCONST | Ops.GEP | Ops.SPECIAL | Ops.UNROLL | Ops.CONTRACT | Ops.SINK | \
            Ops.LINEAR | Ops.PROGRAM | Ops.SOURCE | Ops.BINARY | Ops.INS:
         return None
 
-      case Ops.CAST:
-        # when PTX cases from ptr to non ptr, remove the shape
-        if isinstance(self.src[0].dtype, PtrDType) and not isinstance(self.src[0].dtype, ImageDType) and not isinstance(self.dtype, PtrDType):
-          return None
-
       case Ops.INDEX:
         # non pointer index doesn't have a shape
         if not isinstance(self.dtype, PtrDType): return None
@@ -225,7 +220,7 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
         return self.src[0].shape[len(self.src[1:]):]
 
       # some ops init the shape
-      case Ops.CONST | Ops.VCONST | Ops.DEFINE_VAR | Ops.BIND: return ()
+      case Ops.CONST | Ops.DEFINE_VAR | Ops.BIND: return () if self._device is not None else None
       case Ops.BUFFER: return (self.arg,)
       case Ops.BUFFER_VIEW: return (self.arg[0],)
       case Ops.ENCDEC: return self.arg[0]
@@ -245,8 +240,7 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
       case Ops.BITCAST:
         ps = self.src[0]._shape
         if ps is None: return None
-        if (output_sz:=self.dtype.itemsize) != (input_sz:=self.src[0].dtype.itemsize):
-          return ps[:-1]+(ssimplify((ps[-1]*input_sz) // output_sz),) if len(ps) > 0 else ps
+        if (output_sz:=self.dtype.itemsize) != (input_sz:=self.src[0].dtype.itemsize): return ps[:-1]+(ssimplify((ps[-1]*input_sz) // output_sz),)
         return ps
 
       # TODO: disallow reshape from nothing. tested by TestOpenClip.test_multigpu_clip_score
@@ -374,7 +368,7 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
     return vmin
   def __bool__(self): return self._eval((dtypes.bool,), bool)
   def __int__(self): return self._eval(dtypes.ints, int)
-  def __float__(self): return float(self._eval(dtypes.floats, float))
+  def __float__(self): return self._eval(dtypes.floats, float)
   def substitute(self, dvars:dict[UOp, UOp], name:str|None=None, extra_pm:PatternMatcher|None=None):
     dvars = {k:v for k,v in dvars.items() if k is not v}
     if len(dvars) == 0: return self
@@ -603,6 +597,7 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
     return ret
 
   # in these four, if the shape doesn't change we can return self
+  def forced_reshape(self, arg:tuple[sint, ...]): return self._mop(Ops.RESHAPE, arg, same_shape_noop=False)
   #def reshape(self, arg:tuple[sint, ...]): return self._mop(Ops.RESHAPE, arg, same_shape_noop=True)
   #def expand(self, arg:tuple[sint, ...]): return self._mop(Ops.EXPAND, arg, same_shape_noop=True)
   #def shrink(self, arg:tuple[tuple[sint, sint], ...]): return self._mop(Ops.SHRINK, arg, same_shape_noop=True)
@@ -660,8 +655,7 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
     if self.op in {Ops.CONTIGUOUS, Ops.RESHAPE}: return self.src[0].buffer
     # this buffer can process disk tensors and simple movement ops
     if self is not self.base:
-      from tinygrad.schedule.rangeify import pm_mops
-      from tinygrad.uop.symbolic import symbolic
+      from tinygrad.schedule.rangeify import pm_mops, symbolic
       out = graph_rewrite(self.flatten().index(UOp.range(self.size, 0)), pm_mops+symbolic)
       buf = out.src[0].buffer
       assert isinstance(buf, Buffer), "must be a Buffer for movement ops"
@@ -700,8 +694,7 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
     if self.op not in (Ops.BUFFER, Ops.MSTACK): return None
     # LUNIQUEs are never realized
     if self.op_in_backward_slice_with_self(Ops.LUNIQUE): return None
-    # NOTE: this is used by the JIT to determine which inputs we capture
-    return self.buffer if self.buffer.is_allocated() else None
+    return self.buffer
   @property
   def is_realized(self) -> bool: return self.base.realized is not None
 
@@ -808,7 +801,6 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
     # float has NAN issue and we use explicit NAN in transcendental
     if self.op is Ops.WHERE and dtypes.is_int(self.dtype): return min(self.src[1].vmin, self.src[2].vmin), max(self.src[1].vmax, self.src[2].vmax)
     # NOTE: returned UOp is assumed to be CONST
-    if self.op is Ops.PARAM and len(self.src) >= 4: return self.src[2].arg, self.src[3].arg
     if self.op is Ops.DEFINE_VAR and self.arg: return self.arg[1], self.arg[2]
     if self.op in (Ops.RANGE, Ops.SPECIAL): return 0, (self.src[0]-1).vmax
     if self.op is Ops.BIND: return self.src[0]._min_max # ignore the bound value
@@ -859,11 +851,8 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
 
   # TODO: this should replace placeholder
   @staticmethod
-  def param(slot:int, dtype:DType, shape:tuple[sint, ...]|None=None, device=None, vmin_vmax:tuple[PyConst, PyConst]|None=None, name=None):
-    src: tuple[UOp, ...] = (UOp(Ops.NOOP) if shape is None else shape_to_shape_arg(shape),) + \
-                           (UOp(Ops.NOOP) if device is None else UOp(Ops.DEVICE, arg=device),)
-    if vmin_vmax is not None: src += (UOp.const(dtype, vmin_vmax[0]), UOp.const(dtype.scalar(), vmin_vmax[1]))
-    if name is not None: src += (UOp(Ops.NOOP, arg=name),)
+  def param(slot:int, dtype:DType, shape:tuple[sint, ...]|None=None, device=None):
+    src = (UOp(Ops.NOOP) if shape is None else shape_to_shape_arg(shape),) + (() if device is None else (UOp(Ops.DEVICE, arg=device),))
     return UOp(Ops.PARAM, dtype, src, arg=slot)
 
   def call(self, *srcs:UOp, grad_fxn:Callable|None=None, metadata:tuple[Metadata, ...]=()) -> UOp:
@@ -1239,13 +1228,12 @@ if TRACK_MATCH_STATS or PROFILE:
 SENTINEL: Final[UOp] = cast(UOp, object())
 class BottomUpGate(Exception): pass
 class RewriteContext:
-  def __init__(self, pm, bpm, ctx=None, rewrite_into_calls=False):
+  def __init__(self, pm, bpm, ctx=None):
     self.pm: PatternMatcher|None = pm
     self.bpm: PatternMatcher|None = bpm
     self.bpm_cache: dict[UOp, UOp|None] = {}
     self.ctx = ctx
     self.replace: dict[UOp, UOp] = {}
-    self.rewrite_into_calls = rewrite_into_calls
 
   # no cache needed: pm_rewrite is called at most once per UOp due to the replace dict check in unified_rewrite
   def pm_rewrite(self, x:UOp) -> UOp|None: return unwrap(self.pm).rewrite(x, self.ctx)
@@ -1280,10 +1268,6 @@ class RewriteContext:
             if n in waitlist: stack.extend(waitlist.pop(n))
             continue
         stack.append((n, 1, new_n))
-        # NOTE: CALL is handled as a special case.
-        # The function that is called is not included in the graph_rewrite.
-        # If you want to graph_rewrite a call, you can
-        if new_n.op is Ops.CALL and not self.rewrite_into_calls: self.replace[new_n.src[0]] = new_n.src[0]
         for x in reversed(new_n.src):
           if x in on_stack: continue
           stack.append((x, 0, x))
@@ -1322,10 +1306,22 @@ class RewriteContext:
     return self.replace[root]
 
 @profile_matches
-def graph_rewrite(sink:UOp, pm:PatternMatcher, ctx=None, bottom_up=False, name=None, bpm=None, rewrite_into_calls=False) -> UOp:
-  rewrite_ctx = RewriteContext(pm if not bottom_up else None, pm if bottom_up else bpm, ctx, rewrite_into_calls=rewrite_into_calls)
+def graph_rewrite(sink:UOp, pm:PatternMatcher, ctx=None, bottom_up=False, name=None, bpm=None) -> UOp:
+  rewrite_ctx = RewriteContext(pm if not bottom_up else None, pm if bottom_up else bpm, ctx)
   return rewrite_ctx.unified_rewrite(sink)
 
+@profile_matches
+def graph_rewrite_map(sink:UOp, pm:PatternMatcher, ctx=None, bottom_up=False, name=None, bpm=None,
+                      input_map:dict[UOp, UOp]|None=None, ) -> dict[UOp, UOp]:
+  rewrite_ctx = RewriteContext(pm if not bottom_up else None, pm if bottom_up else bpm, ctx)
+  new_map: dict[UOp, UOp] = {}
+  for k in (list(sink.toposort())[::-1] if bottom_up else sink.toposort()):
+    new_map[k] = v = rewrite_ctx.unified_rewrite(k)
+    if k is not v and k.metadata is not None: all_metadata[v] = tuple(dedup(all_metadata.get(v, ())))+k.metadata
+  if input_map is not None:
+    for k,v in input_map.items(): new_map[k] = new_map.get(v,v)
+  return new_map
+
 def sint_to_uop(x:sint, dtype=dtypes.index) -> UOp: return UOp.const(dtype, x) if isinstance(x, int) else x.cast(dtype)
 
 def select_dtype(u): return (dtypes.long if u.overflows(dtypes.int32) else dtypes.int).vec(u.dtype.count)
@@ -1358,6 +1354,7 @@ _substitute = PatternMatcher([(UPat(tuple(Ops), name="x"), lambda ctx,x: ctx.get
 _remove_all_tags = PatternMatcher([(UPat(GroupOp.All, name="x"), lambda x: x.replace(tag=None) if x.tag is not None else None)])
 
 def gate_kernel_sink(x:UOp) -> bool: return not (x.op is Ops.SINK and isinstance(x.arg, KernelInfo))
+pm_gate_kernel_sink = PatternMatcher([(UPat(Ops.SINK, name="sink"), lambda sink: None if gate_kernel_sink(sink) else panic(BottomUpGate))])
 
 def do_unbind(ctx:dict[Variable, int], x:UOp):
   v,i = x.unbind()
@@ -1448,6 +1445,8 @@ pm_pyrender_extra = PatternMatcher([
   (UPat(Ops.INDEX, src=(UPat(), UPat()), allow_any_len=True, name="x"), lambda ctx,x:
    f"{ctx[x.src[0]]}.index({ctx[x.src[1]]}, "+(f"{ctx[x.src[2]]}, " if len(x.src) > 2 else "")+
     (f"dtype={x.dtype})" if x.src[0].dtype != x.dtype else "ptr=True)") if x.src[0].dtype.base != x.dtype else None),
+  # TODO: fix forced_reshape
+  (UPat(Ops.RESHAPE, name="x"), lambda ctx,x: f"{ctx[x.src[0]]}.forced_reshape({render_marg(ctx,x)})" if x.src[0].shape == x.shape else None),
   (UPat(GroupOp.Movement, name="x"), lambda ctx,x: f"{ctx[x.src[0]]}.{x.op.name.lower()}({render_marg(ctx,x)})"),
   # NOTE: CMPNE doesn't work cause there's no __rne__
   # NOTE: only match CONSTs without UNIQUE (len(src)==1), unique_const needs explicit rendering

tinygrad/uop/spec.py

@@ -58,9 +58,6 @@ shared_spec = PatternMatcher([
 
   # RANGE/SPECIAL define loops, END closes them
   (UPat(Ops.END, src=(UPat(), UPat(Ops.RANGE))), lambda: True),
-
-  # NOOP
-  (UPat(Ops.NOOP), lambda: True)
 ])
 
 # ***** UOp spec in the Tensor graph *****

tinygrad/uop/symbolic.py

@@ -105,7 +105,6 @@ symbolic_simple = propagate_invalid + PatternMatcher([
   (UPat(Ops.BITCAST, name="root", src=(UPat.cvar("c"),)), fold_bitcast),
   # b.cast(a).cast(b) -> b if a preserves all values in b
   (UPat.var('x').cast(name="a").cast(name="b"), lambda x,a,b: x if x.dtype == b.dtype and can_lossless_cast(b.dtype, a.dtype) else None),
-  (UPat.var("x").cast(dtypes.bool), lambda x: x != 0),
   # ** pow **
   (UPat.var("x").alu(Ops.POW, UPat.cvar("c", vec=False)), simplify_pow),
   # positive const ** x
@@ -396,6 +395,9 @@ sym = symbolic+pm_simplify_valid+PatternMatcher([
   # reorder ALU/VECTORIZE
   (UPat(GroupOp.ALU, src=(UPat(Ops.VECTORIZE, src=UPat(name='x')), UPat(Ops.VECTORIZE, src=UPat(name='y'))), name='alu'),
    lambda x,y,alu: UOp(Ops.VECTORIZE, alu.dtype, (UOp(alu.op, alu.dtype.scalar(), (x,y)),)*alu.dtype.count)),
+  # ** self folding **
+  # x!=0 -> (bool)x
+  (UPat.var("x")!=0, lambda x: x.cast(dtypes.bool.vec(x.dtype.count))),
   # ** where **
   # # fold nested where with same condition: in cond.where(t,f), cond.where(a,b)->a in t, ->b in f
   # (UPat.var("cond").where(UPat.var("t"), UPat.var("f")), fold_where_closure),

tinygrad/viz/index.html

@@ -76,7 +76,7 @@
     pointer-events: none;
   }
   label {
-    display: flex;
+    display: inline-flex;
     align-items: center;
     gap: 4px;
     line-height: 1;