Merge branch 'master' into bitcast_spec

skip slow whisper tests
kerenl opts test from nan in llama 8b (#17510 )
2026-08-17 07:18:26 +00:00 · 2026-08-12 19:24:33 -07:00 · 2026-08-12 13:12:37 -07:00 · 2026-08-13 04:07:30 +09:00 · 2026-08-12 08:39:29 -07:00 · 2026-08-12 13:08:53 +03:00
35 changed files with 515 additions and 392 deletions
@@ -521,10 +521,7 @@ jobs:
      - name: Run HCQ2 tests
        run: HCQ_RUNTIME_DEV=PYTHON HCQ2=1 DEV=MOCKKFD+AMD FORWARD_ONLY=1 PYTHONPATH=. python test/test_tiny.py
      - name: Run HCQ2 multi-device tests
-        run: |
-          HCQ_RUNTIME_DEV=PYTHON HCQ2=1 DEV=MOCKKFD+AMD FORWARD_ONLY=1 PYTHONPATH=. python test/unit/test_multitensor.py \
-            TestMultiTensor.test_simple_add TestMultiTensor.test_shard_reduce \
-            TestMultiTensor.test_backward_sum TestMultiTensor.test_matmul_shard_0_0
+        run: HCQ_RUNTIME_DEV=PYTHON HCQ2=1 DEV=MOCKKFD+AMD FORWARD_ONLY=1 PYTHONPATH=. python -m pytest -n=auto test/backend/test_multitensor.py
      - name: Run HCQ2 JIT tests
        run: HCQ_RUNTIME_DEV=PYTHON HCQ2=1 DEV=MOCKKFD+AMD FORWARD_ONLY=1 PYTHONPATH=. python test/unit/test_jit.py
      - name: Run HCQ2 unit tests
@@ -12,7 +12,7 @@ from tinygrad.helpers import Timing, colored, GlobalCounters, profile_marker
 from tinygrad.uop.ops import Ops, UOp
 from extra.models.llama import apply_rotary_emb
 from extra.llama_kernels.rmsnorm import rmsnorm
-from extra.gemm.cdna_asm_gemm import _mx_block_scale, _mx_block_scale_3d, quantize_mxfp8
+from extra.gemm.cdna_asm_gemm import _mx_block_scale, _mx_block_scale_3d, quantize_mxfp8, asm_gemm, can_use_asm_gemm
 from extra.gemm.moe_gemm import grouped_mx_gemm
 from extra.gemm.moe_routing import route, dispatch, combine

@@ -263,7 +263,11 @@ class GPTOSS:
                        w_down=self.w_down[i], w_down_scale=self.w_down_scale[i], w_down_bias=self.w_down_bias[i])
      h, *_ = self.run_layer(h, freqs_cis, mask_full, i % 2 == 0, attn_kwargs, ffn_kwargs, save=save)

-    logits = self.norm(h) @ self.output.T
+    h_normed = self.norm(h)
+    pad = (-self.dim) % 256
+    h_padded, w_padded = h_normed.pad((None, None, (0, pad))), self.output.pad(((0, 0), (0, pad)))
+    if ASM_GEMM and can_use_asm_gemm(h_padded, w_padded.T): logits = asm_gemm(h_padded, w_padded.T)
+    else: logits = h_normed @ self.output.T
    return logits

 def _get_pads(uop:UOp) -> list[UOp]:
@@ -46,7 +46,7 @@ export DATA_SEED=${DATA_SEED:-5760}
 export JITBEAM=${JITBEAM:-3}
 export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=1

-export FAKEDATA=${FAKEDATA:-1} BENCHMARK=${BENCHMARK:-10}
+export FAKEDATA=${FAKEDATA:-$([[ "$DEV" == NULL:* ]] && echo 1 || echo 0)} BENCHMARK=${BENCHMARK:-10}
 if [ -z "$FULL_LAYERS" ]; then
  export LLAMA_LAYERS=${LLAMA_LAYERS:-2}
 fi
@@ -1,8 +1,8 @@
 #!/usr/bin/env bash

 export PYTHONPATH="."
-export PATH="/opt/rocm-7.1.1/bin:$PATH"
-export ROCM_PATH="/opt/rocm-7.1.1"
+export ROCM_PATH=${ROCM_PATH:-/opt/rocm-7.1.1}
+export PATH="$ROCM_PATH/bin:$PATH"
 export DEV=${DEV:-AMD}
 export CHECK_OOB=0
 export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
@@ -297,6 +297,8 @@ class AMDAllocator(HCQAllocator['AMDDevice']):

  def _do_map(self, buf:HCQ2Buffer): return self.dev.iface.map(buf._base if buf._base is not None else buf)

+  def _do_unmap(self, buf:HCQ2Buffer): self.dev.iface.unmap(buf)
+
@dataclass
 class AMDQueueDesc:
  ring: Buffer; read_ptr: Buffer; write_ptr: Buffer; doorbell: Buffer; put_value: Buffer  # noqa: E702
@@ -388,15 +390,24 @@ class KFDIface:
    return hcqbuf

  def free(self, mem):
+    self._unmap(mem)
+    if mem.va_addr: FileIOInterface.munmap(mem.va_addr, mem.size)
+    kfd.AMDKFD_IOC_FREE_MEMORY_OF_GPU(self.kfd, handle=mem.meta.handle)
+
+  def unmap(self, mem):
+    self._unmap(mem)
+    if getattr(mem, '_owns_kfd_handle', False): kfd.AMDKFD_IOC_FREE_MEMORY_OF_GPU(self.kfd, handle=mem.meta.handle)
+
+  def _unmap(self, mem):
    gpus = (ctypes.c_int32 * 1)(self.gpu_id)
    stm = kfd.AMDKFD_IOC_UNMAP_MEMORY_FROM_GPU(self.kfd, handle=mem.meta.handle, device_ids_array_ptr=ctypes.addressof(gpus), n_devices=1)
    assert stm.n_success == 1
-    if mem.owner == self.dev:
-      if mem.va_addr: FileIOInterface.munmap(mem.va_addr, mem.size)
-      kfd.AMDKFD_IOC_FREE_MEMORY_OF_GPU(self.kfd, handle=mem.meta.handle)

  def map(self, mem):
-    if mem.owner is not None and mem.owner._is_cpu(): return self.alloc(mem.size, host=True, cpu_addr=mem.va_addr)
+    if mem.owner is not None and mem.owner._is_cpu():
+      mapped = self.alloc(mem.size, host=True, cpu_addr=mem.va_addr)
+      mapped._owns_kfd_handle = True
+      return mapped

    c_gpus = (ctypes.c_int32 * 1)(self.gpu_id)
    stm = kfd.AMDKFD_IOC_MAP_MEMORY_TO_GPU(self.kfd, handle=mem.meta.handle, device_ids_array_ptr=ctypes.addressof(c_gpus), n_devices=1)
@@ -468,6 +479,7 @@ class PCIIface(PCIIfaceBase):

  def require_profile_mode(self): return True
  def is_wgp_active(self, xcc, se, sa, wgp) -> bool: return True # TODO: account for WGP disablement on some asics.
+  def unmap(self, mem): self.free(mem)

  def _compute_props(self):
    self.ip_versions = self.dev_impl.ip_ver
@@ -0,0 +1,221 @@
+# Runbook: Llama 3 8B Training on DigitalOcean MI350X
+
+## Machine Specs
+- 8x MI350X GPUs (gfx950, device ID 75b0), 288GB VRAM each
+- 2TB RAM, 192 CPUs, 2TB disk
+- ROCm 7.14 at `/opt/rocm` (NOT `/opt/rocm-7.1.1` like the submission scripts assume)
+- Python 3.12
+
+## Phase 1: System Setup
+
+### 1.1 Install packages
+```bash
+apt-get update
+apt-get install -y python3-pip python3-venv git tmux rclone clang
+```
+
+### 1.2 Install Python deps
+```bash
+python3 -m pip install --break-system-packages numpy tqdm wandb tiktoken sentencepiece
+```
+
+### 1.3 Install ROCm dev headers
+The base image has ROCm runtime but NOT the HIP dev headers. Need:
+```bash
+apt-get install -y amdrocm-core-dev
+```
+This installs `hip/hip_runtime.h` at `/opt/rocm/core-7.14/include/hip/hip_runtime.h`.
+The symlink `/opt/rocm/include` → `/opt/rocm/core-7.14/include` makes it available at `/opt/rocm/include/hip/hip_runtime.h`.
+
+### 1.4 Configure ROCm comgr
+ROCm 7.14 ships comgr 3.3 at `/opt/rocm/lib/libamd_comgr.so`. tinygrad's DLL loader needs explicit env vars to find it (it searches for `libcomgr.so*` by default, not `libamd_comgr.so*`). Set these in the run command:
+```bash
+export COMGR_PATH=/opt/rocm/lib/libamd_comgr.so
+export COMGR_3_PATH=/opt/rocm/lib/libamd_comgr.so
+```
+Also add ROCm libs to ldconfig so comgr's shared library dependencies resolve:
+```bash
+cat > /etc/ld.so.conf.d/rocm.conf << 'EOF'
+/opt/rocm/lib
+/opt/rocm/lib/llvm/lib
+/opt/rocm/lib/rocm_sysdeps/lib
+EOF
+ldconfig
+```
+
+### 1.5 Install geohot tmux config
+```bash
+curl -sL https://raw.githubusercontent.com/geohot/configuration/master/.tmux.conf -o ~/.tmux.conf
+```
+
+### 1.6 Reload amdgpu driver
+tinygrad's HCQ backend needs `/dev/kfd` which is created by the amdgpu kernel driver.
+If the driver was unloaded, reload it:
+```bash
+modprobe amdgpu
+ls /dev/kfd  # should exist
+```
+
+## Phase 2: Clone tinygrad
+```bash
+cd /root
+git clone https://github.com/tinygrad/tinygrad.git
+cd tinygrad
+python3 -m pip install --break-system-packages -e .
+```
+
+## Phase 3: Download C4 Dataset
+
+The C4 data is on the MLCommons Cloudflare R2 bucket in Megatron-LM indexed format.
+
+```bash
+rclone config create mlc-training s3 provider=Cloudflare \
+  access_key_id=76ea42eadb867e854061a1806220ee1e \
+  secret_access_key=a53625c4d45e3ca8ac0df8a353ea3a41ffc3292aa25259addd8b7dc5a6ce2936 \
+  endpoint=c2686074cb2caf5cbaf6d134bdba8b47.r2.cloudflarestorage.com
+
+mkdir -p /raid/datasets/c4-8b
+rclone copy mlc-training:mlcommons-training-wg-public/llama3_1/datasets/c4/llama3_1_8b/ /raid/datasets/c4-8b/ -P
+```
+
+Files downloaded (~85GB total, ~6 minutes):
+- `c4-train.en_6_text_document.bin` (79 GB)
+- `c4-train.en_6_text_document.idx` (870 MB)
+- `c4-validation-91205-samples.en_text_document.bin` (159 MB)
+- `c4-validation-91205-samples.en_text_document.idx` (1.8 MB)
+- `LICENSE.txt`, `NOTICE.txt`
+
+## Phase 4: wandb Login
+```bash
+wandb login
+```
+Enter API key from https://wandb.ai/authorize
+
+## Phase 5: Run Training
+
+### 5.1 Smoke test (beam search, 2 layers, real data)
+Always run beam first to validate the pipeline:
+```bash
+cd /root/tinygrad
+COMGR_PATH=/opt/rocm/lib/libamd_comgr.so \
+COMGR_3_PATH=/opt/rocm/lib/libamd_comgr.so \
+CC=/opt/rocm/core-7.14/lib/llvm/bin/clang \
+DEV=AMD:HIP \
+ROCM_PATH=/opt/rocm \
+  bash examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_beam.sh
+```
+
+### 5.2 Full training run
+```bash
+cd /root/tinygrad
+COMGR_PATH=/opt/rocm/lib/libamd_comgr.so \
+COMGR_3_PATH=/opt/rocm/lib/libamd_comgr.so \
+CC=/opt/rocm/core-7.14/lib/llvm/bin/clang \
+DEV=AMD:HIP \
+ROCM_PATH=/opt/rocm \
+WANDB=1 \
+  bash examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_run.sh
+```
+
+## Environment Variable Reference
+
+| Variable | Value | Why |
+|---|---|---|
+| `COMGR_PATH` | `/opt/rocm/lib/libamd_comgr.so` | tinygrad's DLL loader needs explicit path to find comgr 3.3 |
+| `COMGR_3_PATH` | `/opt/rocm/lib/libamd_comgr.so` | comgr 3.x uses a separate `comgr_3` module with its own path var |
+| `CC` | `/opt/rocm/core-7.14/lib/llvm/bin/clang` | System clang doesn't know gfx950; must use ROCm's bundled clang |
+| `DEV` | `AMD:HIP` | Force HIPRenderer (comgr-based) over HIPCCRenderer (hipcc subprocess) |
+| `ROCM_PATH` | `/opt/rocm` | Script defaults to `/opt/rocm-7.1.1` which doesn't exist |
+| `WANDB` | `1` | Enable wandb logging (off by default) |
+
+## Architecture
+
+| Component | Source file |
+|---|---|
+| Model | `examples/mlperf/models/flat_llama.py` — FlatTransformer, FP8 MXFP4 weights, fused QKV, flash attention |
+| Trainer | `examples/mlperf/model_train.py` → `train_llama3()` |
+| Optimizer | `examples/mlperf/optim.py` — GradAccClipAdamW, master weights, FP8 re-quant |
+| LR schedule | `examples/mlperf/lr_schedulers.py` — CosineAnnealingLRWithWarmup |
+| Dataloader | `examples/mlperf/dataloader.py` — Megatron-LM indexed bin format |
+| ASM GEMM | `extra/gemm/cdna_asm_gemm.py` — gfx950 MFMA assembly, MXFP4 |
+| Flash attention | `extra/thunder/amd/fa.py` |
+| Fused kernels | `extra/llama_kernels/` — rmsnorm, silu, quantize, fused_ce |
+| GPU driver | `tinygrad/runtime/ops_amd.py` — HCQ, direct KFD ioctl |
+| Renderer | `tinygrad/renderer/cstyle.py` — HIPRenderer for gfx950 |
+| comgr compiler | `tinygrad/runtime/support/compiler_amd.py` — HIPCompiler using comgr 3.3 |
+
+## Troubleshooting
+
+### `'hip/hip_runtime.h' file not found`
+Install `amdrocm-core-dev`:
+```bash
+apt-get install -y amdrocm-core-dev
+```
+
+### `'gfx950' is not a recognized processor` + LLVM crash
+System clang doesn't know gfx950. Set `CC=/opt/rocm/core-7.14/lib/llvm/bin/clang`.
+
+### `comgr not available: try setting COMGR_PATH?`
+Add ROCm libs to ldconfig and set `COMGR_PATH` and `COMGR_3_PATH`:
+```bash
+# /etc/ld.so.conf.d/rocm.conf should contain /opt/rocm/lib paths
+ldconfig
+```
+
+### `comgr not available: try setting COMGR_3_PATH?`
+comgr 3.x uses a separate module. Set `COMGR_3_PATH=/opt/rocm/lib/libamd_comgr.so` too.
+
+### `No such file or directory: 'clang'`
+Install clang: `apt-get install -y clang` (for CPU compilation).
+For gfx950 HIP compilation, comgr (not clang) is used — ensure the ROCm 7.14 comgr 3.3 is properly loaded via `COMGR_PATH` and `COMGR_3_PATH`.
+
+## Appendix: KVM Virtualization Observations
+
+### Virtualization detection
+```
+$ systemd-detect-virt
+kvm
+$ lspci -nn | grep AMD
+83:00.0 ... Device [1002:75b0]
+```
+CPU flags include `hypervisor`. `dmesg` shows `Hypervisor detected: KVM`.
+
+### PCI device ID
+`lspci -v` shows device ID `0x75b0` and subsystem ID `0x75a0`:
+```
+83:00.0 Processing accelerators: ... Device 75b0
+    Subsystem: ... Device 75a0
+```
+tinygrad's `PCIIface` in `ops_amd.py` and `hive_reset.py` did not list `0x75b0`, so the GPU was not found. Adding `0x75b0` to the device ID list in both files fixes the detection.
+
+### amdgpu driver behavior
+On first boot, amdgpu loaded and bound to all 8 GPUs. On one boot it failed to initialize:
+```
+[  799.780369] amdgpu 0000:83:00.0: Failed to alloc msi vectors
+[  799.781476] amdgpu 0000:83:00.0: sw_init of IP block <vega20_ih> failed -22
+[  799.782724] amdgpu 0000:83:00.0: amdgpu_device_ip_init failed
+[  799.793885] amdgpu 0000:83:00.0: Fatal error during GPU init
+```
+On a subsequent boot, amdgpu initialized successfully (SMU initialized, VRAM ready). After unbinding all 8 GPUs from amdgpu, `rmmod amdgpu` wedged the module (stuck in "Unloading" state in `/proc/modules`), requiring a full VM reboot.
+
+### `/dev/kfd`
+`/dev/kfd` exists when amdgpu is loaded. Opening it returns `OSError: [Errno 22] Invalid argument`.
+
+### VRAM BAR reads all 0xFF
+After amdgpu initializes the GPU and is then unbound, reading the VRAM BAR (via `/sys/bus/pci/devices/0000:83:00.0/resource0`) returns all `0xFF` at all offsets — including the discovery table at `vram_size - 64KB`. tinygrad's `AMDev._run_discovery()` fails with `AssertionError: discovery signatures mismatch`.
+
+A PCI reset (`echo 1 > /sys/bus/pci/devices/0000:83:00.0/reset`) did not change the VRAM contents — still all `0xFF`.
+
+VRAM was also all `0xFF` when read via `/dev/mem` at the BAR physical address (`0xa0000000000`).
+
+### VFIO attempt
+Bound the GPU to `vfio-pci` with `enable_unsafe_noiommu_mode=1`. The GPU bound successfully and `/dev/vfio/noiommu-0` appeared. Running tinygrad with `VFIO=1` still failed with the same `discovery signatures mismatch` — VRAM BAR still reads all `0xFF`.
+
+### No IOMMU in guest
+`dmesg` has no `AMD-Vi` entries. PCI devices have no `iommu_group` symlink.
+
+### No fan control
+No `fan*` or `pwm*` hwmon entries exist. Only `temp*`, `power*`, `freq*` are exposed. GPU temps read 56-63°C, power ~265W per GPU.
+
+### Current status: NOT WORKING
+tinygrad's `PCIIface` finds the GPU (after adding `0x75b0`) but `AMDev._run_discovery()` fails because the VRAM discovery table reads all `0xFF`. This was observed with the GPU unbound from any driver, after PCI reset, and with VFIO bound.
@@ -422,9 +422,8 @@ class TestCustomKernel(unittest.TestCase):
      return Tensor.custom_kernel(y, x, fxn=custom_add_one_kernel)[0]
    GlobalCounters.reset()
    y = run(x[0]).realize()
-    # it's copying the input and the output
-    # TODO: subbuffer usage has runtime specific behavior, this will be fixed after the removal of SLICE.
-    assert_kernel_count(2 if y.device in ("CL", "WEBGPU") else 1)
+    # backends that support contiguous views don't launch extra kernels
+    assert_kernel_count(2 if x[0].uop.contiguous_view() is None else 1)
    self.assertEqual(y.tolist(), [1, 2, 3, 4])

  @Context(DEV="CPU")
@@ -74,6 +74,7 @@ class TestWhisper(unittest.TestCase):
        err
      )

+  @slow
  def test_transcribe_file1(self):
    self.assertEqual(transcribe_file(self.model, self.enc, TEST_FILE_1),  TRANSCRIPTION_1)

@@ -89,6 +90,7 @@ class TestWhisper(unittest.TestCase):
    self.assertEqual(TRANSCRIPTION_1,  transcriptions[0])
    self.assertEqual(TRANSCRIPTION_2,  transcriptions[1])

+  @slow
  def test_transcribe_batch21(self):
    waveforms = [load_file_waveform(TEST_FILE_2), load_file_waveform(TEST_FILE_1)]
    transcriptions = transcribe_waveform(self.model, self.enc, waveforms)
@@ -0,0 +1,45 @@
+import unittest
+from tinygrad import UOp, dtypes
+from tinygrad.uop.ops import shape_to_shape_arg, ParamArg, Ops, AddrSpace
+
+def placeholder(shape, dtype, slot):
+  return UOp(Ops.PARAM, dtype, (shape_to_shape_arg(shape),), arg=ParamArg(slot, AddrSpace.GLOBAL))
+
+class TestBitcastSpec(unittest.TestCase):
+  def test_bitcast_no_shape_change(self):
+    pl = placeholder((10,10), dtypes.int, 0)
+    out = pl.bitcast(dtypes.float)
+    self.assertEqual(out.shape, (10,10))
+
+  def test_bitcast_increase_shape(self):
+    pl = placeholder((10,10), dtypes.int, 0)
+    out = pl.bitcast(dtypes.short)
+    self.assertEqual(out.shape, (10,20))
+
+  def test_bitcast_decrease_shape(self):
+    pl = placeholder((10,10), dtypes.int, 0)
+    out = pl.bitcast(dtypes.long)
+    self.assertEqual(out.shape, (10,5))
+
+  def test_bitcast_remove_ones(self):
+    pl = placeholder((10,2), dtypes.int, 0)
+    out = pl.bitcast(dtypes.long)
+    self.assertEqual(out.shape, (10,1))
+
+  def test_bitcast_remove_ones_full(self):
+    pl = placeholder((2,), dtypes.int, 0)
+    out = pl.bitcast(dtypes.long)
+    self.assertEqual(out.shape, (1,))
+
+  def test_bitcast_add_ones_full(self):
+    pl = placeholder((), dtypes.long, 0)
+    out = pl.bitcast(dtypes.int)
+    self.assertEqual(out.shape, (2,))
+
+  def test_bitcast_add_ones_full_uchar(self):
+    pl = placeholder((), dtypes.long, 0)
+    out = pl.bitcast(dtypes.uchar)
+    self.assertEqual(out.shape, (8,))
+
+if __name__ == '__main__':
+  unittest.main()
@@ -27,10 +27,15 @@ def _make_linear(buffer_lists, copies=None):
    calls.append(UOp(Ops.CALL, src=(src0, *bufs)))
  return UOp(Ops.LINEAR, src=tuple(calls))

+def _get_planned_view(buf:UOp) -> tuple[UOp, int, int]|None:
+  view = buf.src[0] if buf.op is Ops.BITCAST else buf
+  if view.op is not Ops.SHRINK or view.src[0].op is not Ops.BUFFER: return None
+  return (arena:=view.src[0]), view.src[1].val * arena.dtype.itemsize, view.src[2].val * arena.dtype.itemsize
+
 def _get_arena(buf, linear, result):
  for orig_si, new_si in zip(linear.src, result.src):
    for orig, new in zip(orig_si.src[1:], new_si.src[1:]):
-      if orig is buf and new.op is Ops.SLICE: return new.src[0]
+      if orig is buf and (planned:=_get_planned_view(new)) is not None: return planned[0]
  return None

 def check_assign(buffer_lists, copies=None):
@@ -41,8 +46,8 @@ def check_assign(buffer_lists, copies=None):
  replace_map: dict[int, tuple[UOp, int, int]] = {}
  for orig_si, new_si in zip(linear.src, result.src):
    for orig, new in zip(orig_si.src[1:], new_si.src[1:]):
-      if new.op is Ops.SLICE and id(orig) not in replace_map:
-        replace_map[id(orig)] = (new.src[0], new.src[1].val * new.src[0].dtype.itemsize, new.arg * new.dtype.itemsize)
+      if (planned:=_get_planned_view(new)) is not None and id(orig) not in replace_map:
+        replace_map[id(orig)] = planned

  # verify pinned buffers are not planned
  for buf in held_bufs:
@@ -948,6 +948,11 @@ class TestSymbolic(unittest.TestCase):
    self.helper_test_variable(cond.where(u0, u1), 0, 1, "((a<2)!=True)")
    self.helper_test_variable(cond.where(u0, u1).where(u0, u1), 0, 1, "(a<2)")

+  def test_equivalent_const_max(self):
+    x = Variable("x", -10, 10)
+    self.helper_test_variable((x < 0).where(0, x), 0, 10, "x.maximum(0)")
+    self.helper_test_variable((0 < x).where(x, 0), 0, 10, "x.maximum(0)")
+
  def test_where_combine(self):
    cond = Variable("x", 0, 3) < 2
    a = Variable("a", 0, 3)
@@ -239,6 +239,15 @@ class TestKernelOpts(unittest.TestCase):
    helper_linearizer_opt(a.sum().exp(), [[Opt(OptOps.PADTO, 0, 32)],])
    helper_linearizer_opt(a.sum(0).exp(), [[Opt(OptOps.PADTO, 1, 32)],])

+  @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals")
+  @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_shared, "test requires shared")
+  @unittest.expectedFailure
+  def test_padto_group_full_unroll_sum(self):
+    a = Tensor.ones(2, 28, 4096, dtype=dtypes.bfloat16).realize()
+    out = ((a * 0.5).float().square()).sum(axis=(0, 2))
+    opts_to_apply = [Opt(OptOps.GROUPTOP, 1, 256), Opt(OptOps.PADTO, 3, 32), Opt(OptOps.UNROLL, 2, 0), Opt(OptOps.UPCAST, 0, 7)]
+    helper_linearizer_opt(out, [opts_to_apply], check_default_opt=False)
+
  def test_padto_sum(self):
    N = 18
    # NOTE: this setup prevents 17 * 17 contiguous merged into one dimension
@@ -1,5 +1,4 @@
 import unittest
-from unittest.mock import MagicMock
 from tinygrad import Device
 from tinygrad.uop.ops import Ops, UOp
 from tinygrad.dtype import dtypes
@@ -11,36 +10,27 @@ class TestMetalGraph(unittest.TestCase):
    self.MetalGraph = MetalGraph
    self.dev = Device[Device.DEFAULT]

-  def metal_buf(self, offset):
-    buf = MagicMock()
-    if offset > 0:
-      buf.op = Ops.SLICE
-      src = MagicMock()
-      src.dtype = dtypes.uint8
-      buf.src = (src, UOp.const(offset))
-      buf.dtype = dtypes.uint8
-    else:
-      buf.op = Ops.BUFFER
-    buf.device = Device.DEFAULT
-    return buf
+  def metal_buf(self, offset, bitcast=False):
+    size = 4 if bitcast else 1
+    buf = UOp.new_buffer(Device.DEFAULT, offset+size, dtypes.uint8)
+    if offset: buf = buf[offset:offset+size]
+    return buf.bitcast(dtypes.float32) if bitcast else buf

-  def call(self, *bufs):
-    c = MagicMock()
-    c.src = (MagicMock(op=Ops.PROGRAM),) + tuple(bufs)
-    return c
+  def supports_uop(self, *bufs):
+    return self.MetalGraph.supports_uop([self.dev], UOp(Ops.PROGRAM, src=(UOp.sink(),)).call(*bufs))

  def test_supports_uop_normal_offset(self):
-    assert self.MetalGraph.supports_uop([self.dev], self.call(self.metal_buf(0), self.metal_buf(100), self.metal_buf(0xFFFFFFFF))) is True
+    assert self.supports_uop(self.metal_buf(0), self.metal_buf(100), self.metal_buf(0xFFFFFFFF)) is True

  def test_supports_uop_overflow_offset(self):
-    assert self.MetalGraph.supports_uop([self.dev], self.call(self.metal_buf(0), self.metal_buf(0x100000000))) is False
+    assert self.supports_uop(self.metal_buf(0), self.metal_buf(0x100000000)) is False

-  def test_supports_uop_nonmetal_buf(self):
-    # non-SLICE ops should not be checked for offset
-    buf = MagicMock()
-    buf.op = Ops.BUFFER
-    buf.device = Device.DEFAULT
-    self.MetalGraph.supports_uop([self.dev], self.call(buf))
+  def test_supports_uop_non_view_buf(self):
+    assert self.supports_uop(self.metal_buf(0)) is True
+
+  def test_supports_uop_bitcast(self):
+    assert self.supports_uop(self.metal_buf(0xFFFFFFFF, bitcast=True)) is True
+    assert self.supports_uop(self.metal_buf(0x100000000, bitcast=True)) is False

 if __name__ == "__main__":
  unittest.main()
@@ -1,6 +1,7 @@
 import math, functools
 from dataclasses import dataclass
 from tinygrad.dtype import DType, dtypes
+from tinygrad.uop.ops import PatternMatcher, UOp, UPat, Ops

@dataclass(frozen=True)
 class TensorCore: # D = A * B + C, A is (M x K), B is (K x N), C and D are (M x N)
@@ -135,6 +136,42 @@ amd_cdna4 = amd_cdna_1616128 + amd_cdna_161632 + amd_cdna_161616

 def get_amd(arch): return {"gfx942": amd_cdna3, "gfx950": amd_cdna4, "gfx1200": amd_rdna4, "gfx1201": amd_rdna4}.get(arch, amd_rdna3)

+pm_validate_wmma_rdna3 = PatternMatcher([
+  (UPat(Ops.WMMA, name="x", dtype=dtypes.int32), lambda x: x.replace(
+    src=(x.src[0].bitcast(dtypes.uint32), x.src[1].bitcast(dtypes.uint32), x.src[2]))
+    if x.src[0].dtype == dtypes.int8 and x.src[0].max_numel() == 16 else None),
+  (UPat(Ops.WMMA, name="x", dtype=dtypes.half), lambda x: UOp(Ops.STACK, src=tuple(x.replace(
+      src=(x.src[0], x.src[1], UOp(Ops.STACK, src=tuple(x.src[2].index(UOp.const(j//2, dtypes.int16))
+      if j%2 == 0 else UOp.const(0.0, x.src[2].dtype)
+      for j in range(x.max_numel()*2)))),
+      arg=(*x.arg[:4], None)).index(UOp.const(i*2, dtypes.int16))
+      for i in range(x.max_numel()))) if x.max_numel() == 8 else None),
+  (UPat(Ops.WMMA, name="x"), lambda x: x.replace(
+    src=(x.src[0].bitcast(dtypes.uint16), x.src[1].bitcast(dtypes.uint16), x.src[2]))
+    if x.src[0].dtype == dtypes.bfloat16 and x.src[0].max_numel() == 16 else None),
+])
+
+pm_validate_wmma_rdna4 = PatternMatcher([
+  (UPat(Ops.WMMA, name="x", dtype=dtypes.bfloat16), lambda x: x.replace(
+    dtype=dtypes.uint16,
+    src=(x.src[0].bitcast(dtypes.uint16), x.src[1].bitcast(dtypes.uint16), x.src[2].bitcast(dtypes.uint16)))
+      .bitcast(dtypes.bfloat16) if x.max_numel() == 8 and x.src[0].dtype == dtypes.bfloat16 and x.src[0].max_numel() == 8 else None),
+  (UPat(Ops.WMMA, name="x", dtype=dtypes.float),
+    lambda x: x.replace(src=(x.src[0].bitcast(dtypes.uint16), x.src[1].bitcast(dtypes.uint16), x.src[2]))
+    if x.max_numel() == 8 and x.src[0].dtype == dtypes.bfloat16 and x.src[0].max_numel() == 8 else None)
+])
+
+pm_validate_wmma_cdna = PatternMatcher([
+  (UPat(Ops.WMMA, name="x", dtype=dtypes.float),
+    lambda x: x.replace(src=(x.src[0].bitcast(dtypes.uint32), x.src[1].bitcast(dtypes.uint32), x.src[2]))
+    if x.arg[0][2] == 128 and x.src[0].dtype.itemsize <= 8 else None),
+  (UPat(Ops.WMMA, name="x", dtype=dtypes.float),
+    lambda x: x.replace(src=(x.src[0].bitcast(dtypes.uint16), x.src[1].bitcast(dtypes.uint16), x.src[2]))
+    if x.max_numel() == 4 and x.src[0].dtype == dtypes.bfloat16 and x.src[0].max_numel() == 4 else None),
+  (UPat(Ops.WMMA, name="x", dtype=dtypes.float),
+    lambda x: x.replace(src=(x.src[0].bitcast(dtypes.uint64), x.src[1].bitcast(dtypes.uint64), x.src[2]))
+    if x.max_numel() == 4 and x.src[0].dtype in dtypes.fp8_ocp and x.src[0].max_numel() == 8 else None),
+])
 # ***** Apple Metal *****

 metal = [TensorCore(dims=(8,8,8), threads=32, elements_per_thread=(2,2,2), dtype_in=di, dtype_out=do,
@@ -44,8 +44,6 @@ def graph_split_rewrite(linear:UOp, max_batch_size:int=0) -> UOp:
    current_batch, current_batch_devs = [], []

  for si in linear.src:
-    if si.src[0].op is Ops.SLICE: continue
-
    devs = dedup([Device[x] for b in si.src[1:] if b.op is not Ops.BIND for x in (b.device if isinstance(b.device, tuple) else (b.device,))])
    graph_t = graph_class(devs[0]) if devs[0].graph is not None else None

@@ -180,7 +178,7 @@ class CapturedJit(Generic[ReturnType]):
      if call.op is not Ops.CALL: continue
      arg_uops = get_call_arg_uops(call)
      outs, ins = get_call_outs_ins(call)
-      out |= {arg_uops[k] for k in set(outs) - set(ins) if arg_uops[k].op in (Ops.BUFFER, Ops.SLICE)}
+      out |= {b for k in set(outs) - set(ins) if (b:=u if (cv:=(u:=arg_uops[k]).contiguous_view()) is None else cv[0]).op is Ops.BUFFER}
    return out

  def __call__(self, input_uops:list[UOp], var_vals:dict[str, int]) -> ReturnType:
@@ -4,8 +4,8 @@ import time, random, itertools, math, contextlib, weakref, array
 from dataclasses import dataclass, replace, field
 from tinygrad.helpers import colored, DEBUG, GlobalCounters, ansilen, all_int, prod, flatten, Context, getenv, to_tuple
 from tinygrad.helpers import BEAM, size_to_str, time_to_str, VALIDATE_WITH_CPU, PROFILE, ProfilePointEvent, cpu_events, wait_cond
-from tinygrad.uop.ops import Ops, PatternMatcher, UOp, UPat, AxisType, sym_infer, buffers, graph_rewrite
-from tinygrad.device import Device, Buffer, MultiBuffer
+from tinygrad.uop.ops import Ops, PatternMatcher, UOp, UPat, AxisType, sym_infer, graph_rewrite
+from tinygrad.device import Device, Buffer, MultiBuffer, ProfileGraphEntry
 from tinygrad.renderer import Estimates
 from tinygrad.codegen import to_program
 from tinygrad.codegen.opt.postrange import args_from_ast
@@ -17,7 +17,7 @@ def get_call_arg_uops(call:UOp) -> tuple[UOp, ...]: return tuple(s for s in call
 def get_call_outs_ins(call:UOp) -> tuple[tuple[int, ...], tuple[int, ...]]:
  ast = call.src[0]
  if ast.op is Ops.PROGRAM: return tuple(ast.arg.outs), tuple(ast.arg.ins)
-  if ast.op in (Ops.COPY, Ops.SLICE): return (0,), (1,)
+  if ast.op is Ops.COPY: return (0,), (1,)
  if ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "encdec": return (0,), tuple(range(1, len(get_call_arg_uops(call))))
  return (), ()

@@ -27,9 +27,6 @@ def get_call_name(call:UOp, bufs:Sequence[Buffer|UOp], var_vals:dict[str, int]|N

  ast, arg_uops = call.src[0], get_call_arg_uops(call)
  if ast.op is Ops.PROGRAM: return ast.arg.name
-  if ast.op is Ops.SLICE:
-    offset = ast.src[1].val * arg_uops[1].dtype.itemsize
-    return colored(f"view {_uop_sz_to_str(arg_uops[0]):>10} @ {offset:<10d}", "yellow")
  if ast.op is Ops.COPY: return colored(f"copy {_uop_sz_to_str(arg_uops[0]):>10}, {_dev_str(bufs[0]):>7s} <- {_dev_str(bufs[1]):7s}", "yellow")
  if ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "encdec": return colored(f"enc/dec {_uop_sz_to_str(arg_uops[0])}", "yellow")
  if ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "graph": return colored(f"batched {len(ast.src[0].src)}", "cyan")
@@ -140,7 +137,7 @@ class ExecContext:
  cache: bool = True

 def _resolve(b:UOp, inputs:tuple[UOp, ...]) -> UOp:
-  if b.op in (Ops.SLICE, Ops.MSELECT) and b.src[0].op is Ops.PARAM: return b.replace(src=(inputs[b.src[0].arg.slot], *b.src[1:]))
+  if b.op in (Ops.MSELECT, Ops.SHRINK) and b.src[0].op is Ops.PARAM: return b.replace(src=(inputs[b.src[0].arg.slot], *b.src[1:]))
  if b.op is Ops.MSTACK: return b.replace(src=tuple(_resolve(x, inputs) for x in b.src))
  return inputs[b.arg.slot] if b.op is Ops.PARAM else b
 def resolve_params(call:UOp, inputs:tuple[UOp, ...]) -> list[UOp]: return [_resolve(b, inputs) for b in get_call_arg_uops(call)]
@@ -154,13 +151,6 @@ def unwrap_multi(call:UOp, resolved:list[UOp]) -> Iterator[tuple[list[Buffer], d
                   for x in call.src[0].toposort())
    for j, per_dev in enumerate(zip(*[cast(MultiBuffer, b).bufs for b in bufs])): yield list(per_dev), {"_device_num": j} if has_dnum else {}

-def exec_view(ctx:ExecContext, call:UOp, ast:UOp) -> float|None:
-  resolved = resolve_params(call, ctx.input_uops)
-  bufs = [cast(Buffer, b.buffer) for b in resolved]
-  bv = bufs[1].view(resolved[0].max_numel(), ast.dtype, ast.src[1].val*bufs[1].dtype.itemsize)
-  with track_stats(ctx, call, bv.device, [bv, bufs[1]], ctx.var_vals): buffers[resolved[0]] = bv
-  return None
-
 def exec_copy(ctx:ExecContext, call:UOp, ast:UOp) -> float|None:
  for bufs, device_vars in unwrap_multi(call, resolve_params(call, ctx.input_uops)):
    dest, src = bufs[0].ensure_allocated(), bufs[1].ensure_allocated()
@@ -210,27 +200,30 @@ def exec_graph(ctx:ExecContext, call:UOp, ast:UOp) -> float|None:
  return t[0]

 def exec_hcq(ctx:ExecContext, call:UOp, ast:UOp) -> float|None:
-  if (inputs:=call.arg.aux.inputs) is not None:
+  if (info:=call.arg.aux).inputs is not None:
    bufs = [_resolve(ctx.input_uops[i], ctx.input_uops).buffer for i in call.arg.aux.input_idxs]
-    table = call.src[1+inputs].buffer
+    table = call.src[1+info.inputs].buffer
    for j,dev in enumerate(call.arg.aux.device):
      addrs = array.array('Q', [(b.bufs[j] if isinstance(b, MultiBuffer) else b).get_buf(dev).va_addr for b in bufs])
      mv = (table.bufs[j] if isinstance(table, MultiBuffer) else table).ensure_allocated()._buf.cpu_view().view(fmt='Q')
      wait_cond(lambda: mv[0], value=0, timeout_ms=ctx.timeout or getenv("HCQDEV_WAIT_TIMEOUT_MS", 30000), msg=f"{dev} hang detected")
      mv[:len(addrs)] = addrs

-  exec_kernel(replace(ctx, update_stats=False), call, ast)
+  exec_kernel(replace(ctx, update_stats=DEBUG>=3), call, ast)

-  tms:list[float|None] = []
-  for e in (aux:=call.arg.aux).prof: cast(Any, Device[e.device]).prof_ents[e.st_id] = e
-  for d in [cast(Any, Device[x]) for x in aux.device]:
-    with track_stats(ctx, call, d.device, [], ctx.var_vals) as et:
-      if ctx.wait:
-        d.synchronize(timeout=ctx.timeout)
-        ts = [d.signal(i)._buf.cpu_view().view(fmt='Q')[0] for e in aux.prof if e.device == d.device for i in (e.st_id, e.en_id)]
-        if ts: et[0] = float(max(ts)-min(ts))/d.timestamp_divider/1e6
-      tms += et
-  return tms[0]
+  tms = []
+  for devices,name,estimates,prof in info.kernels:
+    for device in devices:
+      tm = None
+      if prof:
+        (d:=cast(Any, Device[device])).prof_ents[prof[0]] = ProfileGraphEntry(device, name, *prof)
+        if ctx.wait:
+          d.synchronize(timeout=ctx.timeout)
+          st, en = (d.signal(x)._buf.cpu_view().view(fmt='Q')[0] for x in prof)
+          tms.append(tm:=float(en-st)/d.timestamp_divider/1e6)
+      stat_call = call.replace(arg=replace(call.arg, name=name, aux=replace(info, estimates=estimates, kernels=())))
+      with track_stats(ctx, stat_call, device, [], ctx.var_vals) as et: et[0] = tm
+  return max(tms) if tms else None

 # flatten LINEAR-in-LINEAR: any nested LINEAR child gets inlined into its parent's src
 pm_flatten_linear = PatternMatcher([
@@ -261,7 +254,6 @@ pm_optimize_local_size = PatternMatcher([
 ])

 pm_exec = PatternMatcher([
-  (UPat(Ops.CALL, src=(UPat(Ops.SLICE, name="ast"),), name="call", allow_any_len=True), exec_view),
  (UPat(Ops.CALL, src=(UPat(Ops.COPY, name="ast"),), name="call", allow_any_len=True), exec_copy),
  (UPat(Ops.CALL, src=(UPat(Ops.PROGRAM, name="ast"),), name="call", allow_any_len=True), exec_kernel),
  (UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION, arg="encdec", name="ast"),), name="call", allow_any_len=True), exec_encdec),
@@ -276,7 +268,7 @@ def compile_linear(linear:UOp, beam:int|None=None, validate=False, input_uops:li
  if validate: linear = graph_rewrite(linear, pm_validate, name="validate", walk=True)
  if (beam_val:=BEAM.value if beam is None else beam) >= 1: linear = graph_rewrite(linear, pm_beam, ctx=beam_val, walk=True)
  linear = graph_rewrite(linear, pm_compile, name="precompile kernels", walk=True)
-  if getenv("HCQ2"): linear = hcq_compile(linear, input_uops, bool(PROFILE) if profile is None else profile)
+  if getenv("HCQ2"): linear = hcq_compile(linear, input_uops, bool(PROFILE or DEBUG >= 2) if profile is None else profile)
  return graph_rewrite(linear, pm_optimize_local_size, name="optimize local size", walk=True)

 def link_linear(linear:UOp, cache=True) -> UOp: return hcq_link(linear, cache=cache) if getenv("HCQ2") else linear
@@ -1,9 +1,9 @@
 from __future__ import annotations
 import time
 START_TIME = time.perf_counter()
-import os, functools, platform, re, contextlib, operator, hashlib, pickle, sqlite3, tempfile, pathlib, string, ctypes, sys, gzip, getpass, gc
+import os, functools, re, contextlib, operator, hashlib, pickle, sqlite3, tempfile, pathlib, string, ctypes, sys, gzip, getpass, gc
 from collections import defaultdict
-import subprocess, shutil, math, types, copyreg, inspect, importlib, decimal, itertools, difflib
+import shutil, math, types, copyreg, inspect, importlib, decimal, itertools, difflib
 from dataclasses import dataclass, field, replace
 from typing import ClassVar, Iterable, Any, TypeVar, Callable, Sequence, TypeGuard, Iterator, Generic, Generator, cast, overload

@@ -13,8 +13,7 @@ U = TypeVar("U")
 def prod(x:Iterable[T]) -> T|int: return functools.reduce(operator.mul, x, 1)

 # NOTE: helpers is not allowed to import from anything else in tinygrad
-OSX, WIN = platform.system() == "Darwin", sys.platform == "win32"
-ARCH_X86 = any(x in platform.processor() for x in ("Intel", "i386", "x86_64"))
+OSX, WIN = sys.platform == "darwin", sys.platform == "win32"
 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
@@ -231,7 +230,7 @@ class _DEV(ContextVar):

 DEV, DEBUG, BEAM, NOOPT = _DEV("DEV", ""), ContextVar("DEBUG", 0), ContextVar("BEAM", 0), ContextVar("NOOPT", 0)
 IMAGE, FLOAT16, OPENPILOT_HACKS = ContextVar("IMAGE", 0), ContextVar("FLOAT16", 0), ContextVar("OPENPILOT_HACKS", 0)
-JIT, JIT_BATCH_SIZE = ContextVar("JIT", 2 if OSX and ARCH_X86 else 1), ContextVar("JIT_BATCH_SIZE", 32)
+JIT, JIT_BATCH_SIZE = ContextVar("JIT", 1), ContextVar("JIT_BATCH_SIZE", 32)
 CHUNK_SIZE = 2**20  # TinyFS content-addressed store: blob chunk + hash-tree node granularity
 WINO, CAPTURING, TRACEMETA, NO_COLOR = ContextVar("WINO", 0), ContextVar("CAPTURING", 1), ContextVar("TRACEMETA", 1), ContextVar("NO_COLOR", 0)
 TRAINING = ContextVar("TRAINING", 0)
@@ -454,9 +453,9 @@ def _ensure_downloads_dir() -> pathlib.Path:
  if pathlib.Path("/etc/tinybox-release").is_file():
    # try creating dir with sudo
    if not (downloads_dir := pathlib.Path("/raid/downloads")).exists():
-      subprocess.run(["sudo", "mkdir", "-p", downloads_dir], check=True)
-      subprocess.run(["sudo", "chown", "tiny:root", downloads_dir], check=True)
-      subprocess.run(["sudo", "chmod", "775", downloads_dir], check=True)
+      system(f"sudo mkdir -p {downloads_dir}")
+      system(f"sudo chown tiny:root {downloads_dir}")
+      system(f"sudo chmod 775 {downloads_dir}")
    return downloads_dir
  return pathlib.Path(cache_dir) / "downloads"

@@ -497,6 +496,7 @@ def fetch_fw(path:str, name:str, sha256:str) -> bytes:
 # *** Exec helpers

 def system(cmd:str, **kwargs) -> str:
+  import subprocess
  st = time.perf_counter()
  try: ret = subprocess.check_output(cmd.split(), stderr=subprocess.STDOUT, **kwargs).decode().strip()
  except subprocess.CalledProcessError as e:
@@ -1,4 +1,4 @@
-import json, math, pathlib, zipfile, pickle, tarfile, struct, functools, io, zlib
+import json, math, pathlib, struct, functools, io, zlib
 from collections import OrderedDict
 from typing import Any, Callable, BinaryIO, Iterable, cast
 from tinygrad.tensor import Tensor
@@ -219,6 +219,7 @@ def load_state_dict(model, state_dict:dict[str, Tensor], strict=True, verbose=Tr

@accept_filename
 def zip_extract(t: Tensor) -> dict[str, Tensor]:
+  import zipfile
  files: dict[str, Tensor] = {}
  with zipfile.ZipFile(TensorIO(t), "r") as myzip:
    # sadly, the extra length needs to be read from the local header of each file.
@@ -249,6 +250,7 @@ def tar_extract(t: Tensor) -> dict[str, Tensor]:
  tensors = nn.state.tar_extract(Tensor(pathlib.Path("archive.tar")))
  ```
  """
+  import tarfile
  with tarfile.open(fileobj=TensorIO(t), mode="r") as tar:
    return {member.name:t[member.offset_data:member.offset_data+member.size] for member in tar if member.type == tarfile.REGTYPE}

@@ -303,6 +305,7 @@ def torch_load(t:Tensor) -> dict[str, Tensor]:
               "FloatTensor": None, "Parameter": Parameter}
  whitelist = {"torch", "collections", "numpy", "_codecs"}  # NOTE: this is not for security, only speed
  class Dummy: pass
+  import pickle, zipfile, tarfile
  class TorchPickle(pickle.Unpickler):
    def find_class(self, module, name):
      module_root = module.split(".")[0]
@@ -279,43 +279,9 @@ exit: %packed = phi i32 [%packed_bf8, %do_bf8], [%packed_fp8, %do_fp8]\n  %trunc
      (UPat(Ops.WMMA, name="wmma"), lambda ctx, wmma, rdna4=AMDLLVMRenderer.is_rdna4(target.arch), cdna=self.is_cdna:
        render_wmma_amd(ctx, wmma, cdna, rdna4))
    ])
-    if self.is_cdna:
-      self.extra_matcher += PatternMatcher([
-        (UPat(Ops.WMMA, name="x", dtype=dtypes.float),
-          lambda x: x.replace(src=(x.src[0].bitcast(dtypes.uint32), x.src[1].bitcast(dtypes.uint32), x.src[2]))
-          if x.arg[0][2] == 128 and x.src[0].dtype.itemsize <= 8 else None),
-        (UPat(Ops.WMMA, name="x", dtype=dtypes.float),
-          lambda x: x.replace(src=(x.src[0].bitcast(dtypes.uint16), x.src[1].bitcast(dtypes.uint16), x.src[2]))
-          if x.max_numel() == 4 and x.src[0].dtype == dtypes.bfloat16 and x.src[0].max_numel() == 4 else None),
-        (UPat(Ops.WMMA, name="x", dtype=dtypes.float),
-          lambda x: x.replace(src=(x.src[0].bitcast(dtypes.uint64), x.src[1].bitcast(dtypes.uint64), x.src[2]))
-          if x.max_numel() == 4 and x.src[0].dtype in dtypes.fp8_ocp and x.src[0].max_numel() == 8 else None),
-      ])
-    if target.arch in {"gfx1100", "gfx1151"}:
-      self.extra_matcher += PatternMatcher([
-        (UPat(Ops.WMMA, name="x", dtype=dtypes.int32), lambda x: x.replace(
-          src=(x.src[0].bitcast(dtypes.uint32), x.src[1].bitcast(dtypes.uint32), x.src[2]))
-          if x.src[0].dtype == dtypes.int8 and x.src[0].max_numel() == 16 else None),
-        (UPat(Ops.WMMA, name="x", dtype=dtypes.half), lambda x: UOp(Ops.STACK, src=tuple(x.replace(
-          src=(x.src[0], x.src[1], UOp(Ops.STACK, src=tuple(x.src[2].index(UOp.const(j//2, dtypes.int16))
-            if j%2 == 0 else UOp.const(0.0, x.src[2].dtype)
-            for j in range(x.max_numel()*2)))),
-          arg=(*x.arg[:4], None)).index(UOp.const(i*2, dtypes.int16))
-          for i in range(x.max_numel()))) if x.max_numel() == 8 else None),
-        (UPat(Ops.WMMA, name="x"), lambda x: x.replace(
-          src=(x.src[0].bitcast(dtypes.uint16), x.src[1].bitcast(dtypes.uint16), x.src[2]))
-          if x.src[0].dtype == dtypes.bfloat16 and x.src[0].max_numel() == 16 else None),
-      ])
-    if target.arch in {"gfx1200", "gfx1201"}:
-      self.extra_matcher += PatternMatcher([
-        (UPat(Ops.WMMA, name="x", dtype=dtypes.bfloat16), lambda x: x.replace(
-          dtype=dtypes.uint16,
-          src=(x.src[0].bitcast(dtypes.uint16), x.src[1].bitcast(dtypes.uint16), x.src[2].bitcast(dtypes.uint16)))
-            .bitcast(dtypes.bfloat16) if x.max_numel() == 8 and x.src[0].dtype == dtypes.bfloat16 and x.src[0].max_numel() == 8 else None),
-        (UPat(Ops.WMMA, name="x", dtype=dtypes.float),
-          lambda x: x.replace(src=(x.src[0].bitcast(dtypes.uint16), x.src[1].bitcast(dtypes.uint16), x.src[2]))
-          if x.max_numel() == 8 and x.src[0].dtype == dtypes.bfloat16 and x.src[0].max_numel() == 8 else None)
-      ])
+    if self.is_cdna: self.extra_matcher += tc.pm_validate_wmma_cdna
+    if target.arch in {"gfx1100", "gfx1151"}: self.extra_matcher += tc.pm_validate_wmma_rdna3
+    if target.arch in {"gfx1200", "gfx1201"}: self.extra_matcher += tc.pm_validate_wmma_rdna4

  def supported_dtypes(self): return {d for d in super().supported_dtypes()
                                      if (d not in dtypes.fp8_ocp or self.target.arch == "gfx950") and d not in dtypes.fp8_fnuz}
@@ -1,4 +1,4 @@
-import glob, importlib, os, pathlib, shutil, subprocess, tarfile, tempfile
+import glob, importlib, os, pathlib, subprocess
 from tinygrad.helpers import fetch, flatten, system, getenv

 root = (here:=pathlib.Path(__file__).parent).parents[2]
@@ -31,6 +31,7 @@ def load(name, files, **kwargs):
  if not (f:=(root/(path:=kwargs.pop("path", __name__)).replace('.','/')/f"{name}.py")).exists() or getenv('REGEN'):
    files, kwargs['args'] = files() if callable(files) else files, args() if callable(args:=kwargs.get('args', [])) else args
    if (srcs:=kwargs.pop('srcs', None)):
+      import tempfile, tarfile
      srcpath = (td:=tempfile.TemporaryDirectory(f"autogen-src-{name.replace('/','-')}")).name + "/"
      for src in (srcs if isinstance(srcs, list) else [srcs]):
        if 'tar' in src:
@@ -157,7 +158,7 @@ def __getattr__(nm):
          *[f"python3 src/compiler/nir/nir_{s}_h.py --outdir gen" for s in ["intrinsics", "intrinsics_indices"]]]), cwd=path, shell=True, check=True),
  srcs="https://gitlab.freedesktop.org/mesa/mesa/-/archive/mesa-25.2.7/mesa-25.2.7.tar.gz",
  dll=f"'tinymesa_cpu' if DEV.renderer == 'LVP' else 'tinymesa', {tinymesa_path}, emsg='pip install tinymesa==25.2.7.2'",
-  prolog=["from tinygrad.helpers import DEV", "import gzip, base64, platform, sysconfig, os"],
+  prolog=["from tinygrad.helpers import DEV", "import gzip, base64, sysconfig, os"],
  epilog=lambda path: [system(f"{root}/extra/mesa/lvp_nir_options.sh {path}")])
    case "libclang":
      return load("libclang",
@@ -5,7 +5,7 @@ from typing import Literal, TypeAlias
 from tinygrad.runtime.support.c import _IO, _IOW, _IOR, _IOWR
 from tinygrad.runtime.support import c
 from tinygrad.helpers import DEV
-import gzip, base64, platform, sysconfig, os
+import gzip, base64, sysconfig, os
 dll = c.DLL('mesa', 'tinymesa_cpu' if DEV.renderer == 'LVP' else 'tinymesa', os.path.join(sysconfig.get_paths()['platlib'], 'tinymesa'), emsg='pip install tinymesa==25.2.7.2')
 class struct_u_printf_info(c.Struct): pass
 u_printf_info: TypeAlias = struct_u_printf_info
@@ -6,7 +6,6 @@ from tinygrad.device import Buffer, BufferSpec, Compiled, Device, MultiBuffer, P
 from tinygrad.dtype import dtypes
 from tinygrad.uop.ops import UOp, Ops, Variable
 from tinygrad.engine.jit import GraphRunner, MultiGraphRunner
-from tinygrad.runtime.ops_rdma import RDMACopyQueue

 class HCQGraph(MultiGraphRunner):
  def __init__(self, *args, **kwargs):
@@ -50,7 +49,7 @@ class HCQGraph(MultiGraphRunner):

    self.comp_queues: dict[HCQCompiled, HWQueue] = {dev: unwrap(dev.hw_compute_queue_t)() for dev in self.devices}
    self.copy_queues: dict[tuple[HCQCompiled, int], HWQueue] = {} # lazy allocation, keyed by (device, queue_idx)
-    self.rdma_queues: dict[tuple[HCQCompiled, HCQCompiled], RDMACopyQueue] = {} # lazy allocation, keyed by device pair
+    self.rdma_queues: dict[tuple[HCQCompiled, HCQCompiled], "RDMACopyQueue"] = {} # lazy allocation, keyed by device pair
    self.num_copy_queues: int = getenv("HCQ_NUM_SDMA", min(len(self.devices), 8) if ALL2ALL >= 1 else 1)
    self.num_rdma_ops: dict[tuple[HCQCompiled, HCQCompiled], int] = collections.defaultdict(int)

@@ -104,6 +103,7 @@ class HCQGraph(MultiGraphRunner):
      elif is_rdma:
        enqueue_queue = self.comp_queues[enqueue_dev]
        rdma_key = (cast(HCQCompiled, Device[bufs[0].device]).rdma_dev(), enqueue_dev.rdma_dev())
+        from tinygrad.runtime.ops_rdma import RDMACopyQueue
        self.rdma_queues.setdefault(rdma_key, RDMACopyQueue(enqueue_dev.rdma_dev()))
      else:
        assert (enqueue_dev.hw_copy_queue_t is not None), "device must implement a copy queue"
@@ -113,5 +113,6 @@ class MetalGraph(GraphRunner):
  @staticmethod
  def supports_uop(batch_devs, new_call:UOp) -> bool:
    # Metal ICB replay encodes offsets as uint32; reject if any Metal buffer offset exceeds 32-bit range.
-    if any(b.op is Ops.SLICE and b.src[1].val * b.src[0].dtype.itemsize > 0xFFFFFFFF for b in new_call.src[1:]): return False
+    for shrink in [s for src in new_call.src[1:] if (s:=src.src[0] if src.op is Ops.BITCAST else src).op is Ops.SHRINK]:
+      if shrink.src[1].val * shrink.src[0].dtype.itemsize > 0xFFFFFFFF: return False
    return GraphRunner.supports_uop(batch_devs, new_call)
@@ -1,5 +1,5 @@
 from __future__ import annotations
-import platform, sys, os, ctypes, functools, mmap, threading, array, itertools
+import platform, sys, os, ctypes, ctypes.util, functools, mmap, threading, array, itertools
 from dataclasses import replace
 from typing import cast
 from tinygrad.helpers import to_mv, OSX, WIN, Context, mv_address, suppress_finalizing, unwrap, data64_le, partition
@@ -33,13 +33,13 @@ class HCQInfo:

  input_idxs:tuple[int, ...] = () # indexes into input_uops used by this call
  inputs:int|None = None
-  prof:tuple[ProfileGraphEntry, ...] = () # st_id/en_id are timestamp signal slots until collect
+  kernels:tuple[tuple[tuple[str, ...], str, Estimates, tuple[int, ...]], ...] = ()

 def all_devices_in(d:Any, c:frozenset[str]) -> bool: return {x.split(":")[0] for x in to_tuple(d)} <= c

 def unwrap_mstack(u):
  if u.op is Ops.MSTACK: return tuple(x for s in u.src for x in unwrap_mstack(s))
-  return unwrap_mstack(u.src[0]) if u.op in {Ops.MSELECT, Ops.SLICE} else (u,)
+  return unwrap_mstack(u.src[0]) if u.op is Ops.MSELECT else (u,)

 def is_value_known_at_link(val:UOp) -> bool:
  runtime_reads = [u for u in val.toposort() if u.op in (Ops.LOAD, Ops.INDEX)]
@@ -141,14 +141,14 @@ def _build_wait_cmds(slots:dict[str, int], dep_lanes:list[tuple[tuple, int, int]

  # opt2: keep latest dep per (dep device, queue, cur lane)
  latest = {((dep[0][dlane], dep[1]), lane): (dep, dlane) for dep, dlane, lane in sorted(dep_lanes, key=lambda x: x[0][2])}
-  deps:dict[tuple, list[int|None]] = collections.defaultdict(lambda: [None]*len(devices))
-  for (_, lane), (dep, dlane) in latest.items(): deps[dep][lane] = dlane
+  deps:dict[tuple, dict[int, list[int]]] = collections.defaultdict(lambda: collections.defaultdict(list))
+  for (_, lane), (dep, dlane) in latest.items(): deps[dep][lane].append(dlane)

  waits = []
-  for (ddevs, dqueue, dtag), lanes in deps.items():
-    sig = UOp.mstack(*[make_signal(d, tag="sentinel_signal") if dl is None else make_signal(ddevs[dl], slots[dqueue])
-                       for dl, d in zip(lanes, devices)])
-    waits.append(UOp(Ops.INS, arg="wait", src=(sig, UOp.const(dtag + 1, dtypes.uint64))))
+  for (ddevs, dqueue, dtag), by_lane in deps.items():
+    for ls in itertools.zip_longest(*(by_lane[lane] for lane in range(len(devices)))):
+      s = UOp.mstack(*[make_signal(d, tag="sentinel_signal") if dl is None else make_signal(ddevs[dl], slots[dqueue]) for dl, d in zip(ls, devices)])
+      waits.append(UOp(Ops.INS, arg="wait", src=(s, UOp.const(dtag + 1, dtypes.uint64))))
  return waits, {dtag for _, _, dtag in deps}

 def _build_finalizers(batch:list[tuple[UOp, tuple[str, ...]]], batch_info:list[tuple[tuple[str, ...], str]],
@@ -199,10 +199,10 @@ def _finalize_batch(batch:list[tuple[UOp, tuple[str, ...]]], profile:bool) -> li
    signal_tags |= cur_signal_tags

  # build fences and finalizers
-  fences, finalizers, finalizer_signal_tags = _build_finalizers(batch, batch_info, deps_tracker, slots)
+  fences, fins, finalizer_signal_tags = _build_finalizers(batch, batch_info, deps_tracker, slots)
  signal_tags |= finalizer_signal_tags

-  src, prof = [], []
+  src, kerns = [], []
  for tag, ((call, _), (devices, queue), q) in enumerate(zip(batch, batch_info, call_waits)):
    # first queue use, sync prior device work with the device timeline
    if batch_info.index((devices, queue)) == tag:
@@ -212,18 +212,18 @@ def _finalize_batch(batch:list[tuple[UOp, tuple[str, ...]]], profile:bool) -> li
    # and make hcq call
    name, info = get_call_name(call, get_call_arg_uops(call)), HCQInfo(devices, estimate_uop(call))
    ts_ids = [next(UOp.unique_num) for _ in range(2)] if profile else []
-    prof += [ProfileGraphEntry(d, name, *ts_ids) for d in devices if ts_ids]
+    kerns.append((devices, name, info.estimates, tuple(ts_ids)))

    ts_ins = [UOp(Ops.INS, arg="timestamp", src=(make_signal(devices, s),)) for s in ts_ids]
    q += ts_ins[:1] + [call.replace(arg=replace(call.arg, aux=info))] + ts_ins[1:]

    # signal the queue if someone waits for us
    if tag in signal_tags: q += [UOp(Ops.INS, arg="store", src=(make_signal(devices, slots[queue]), UOp.const(tag + 1, dtypes.uint64)))]
-    src.append(make_call(name, make_submit(*q, devs=devices, queue=queue).sink(), info))
+    src.append(make_call(f"submit {name}", make_submit(*q, devs=devices, queue=queue).sink(), info))

  # append batch timestamps to finalizers
-  finalizers = [f.replace(arg=replace(f.arg, aux=replace(a:=f.arg.aux, prof=tuple(e for e in prof if e.device in a.device)))) for f in finalizers]
-  return fences + src + finalizers
+  fins = [f.replace(arg=replace(f.arg, aux=replace(a:=f.arg.aux, kernels=tuple(x for x in kerns if set(x[0]) & set(a.device))))) for f in fins]
+  return fences + src + fins

 def sched_hcq_batches(l:UOp, profile:bool) -> UOp:
  srcs:list[UOp] = []
@@ -377,14 +377,15 @@ pm_replace_params = PatternMatcher([

 # *****************

-def resolve_getaddr_slice(bv:UOp, g:UOp) -> UOp:
+def resolve_getaddr_view(bv:UOp, g:UOp) -> UOp:
  base = bv.src[0].after(*g.src[0].src[1:] if g.src[0].op is Ops.AFTER else ())
-  itemsize = bv.src[0].dtype.itemsize if bv.src[0].without_after.op in (Ops.BUFFER, Ops.SLICE, Ops.MSTACK, Ops.MSELECT) else bv.dtype.itemsize
+  if bv.op is Ops.BITCAST: return UOp(Ops.GETADDR, src=(base,), arg=g.arg)
+  itemsize = bv.src[0].dtype.itemsize if bv.src[0].without_after.op in (Ops.BUFFER, Ops.MSTACK, Ops.MSELECT) else bv.dtype.itemsize
  return UOp(Ops.GETADDR, src=(base,), arg=g.arg) + UOp.const(bv.src[1].val * itemsize, dtypes.uint64)

 pm_early_simplify = PatternMatcher([
-  (UPat(Ops.GETADDR, src=(UPat.any(sl:=UPat(Ops.SLICE, name="bv"), sl.after(allow_any_len=True)),), name="g"), resolve_getaddr_slice),
-  (UPat(Ops.INDEX, src=(UPat(Ops.SLICE, name="bv"),), allow_any_len=True, name="x"),
+  (UPat(Ops.GETADDR, src=(UPat((Ops.SHRINK, Ops.BITCAST), name="bv").or_after(),), name="g"), resolve_getaddr_view),
+  (UPat(Ops.INDEX, src=(UPat(Ops.SHRINK, name="bv"),), allow_any_len=True, name="x"),
   lambda bv,x: x.replace(src=(bv.src[0], x.src[1] + bv.src[1].cast(x.src[1].dtype), *x.src[2:]))),
 ])

@@ -402,7 +403,7 @@ def pack_hcq_placeholders(call:UOp) -> UOp|None:
      sizes[b.tag] = offs[b] + b.max_numel()
  counts = collections.Counter(b.tag for b in bufs)
  bases = {b.tag:UOp.placeholder((sizes[b.tag],), b.dtype, next(UOp.unique_num), device=b.device).rtag(b.tag) for b in bufs if counts[b.tag] > 1}
-  subs = {b:UOp(Ops.SLICE, b.dtype, (bases[b.tag], UOp.const(offs.get(b, 0))), b.max_numel()) for b in bufs if b.tag in bases}
+  subs = {b:bases[b.tag][(off:=offs.get(b, 0)):off+b.max_numel()] for b in bufs if b.tag in bases}
  return call.replace(src=(call.src[0].substitute(subs, walk=True), *call.src[1:])) if subs else None
 pm_pack_placeholders = PatternMatcher([
  (UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION, arg="hcq"),), name="call", allow_any_len=True), pack_hcq_placeholders)])
@@ -493,7 +494,7 @@ pm_resolve_patches = PatternMatcher([
  (UPat(name="buf").index(UPat(Ops.RANGE), allow_any_len=True)
    .store(UPat.any(UPat(Ops.BINARY, name="blob"), UPat(Ops.BINARY, name="blob").bitcast()).index(UPat(Ops.RANGE), allow_any_len=True).load())
    .end(UPat(Ops.RANGE)), fold_binary),
-  (UPat({Ops.BUFFER, Ops.SLICE, Ops.MSTACK}, name="buf").index(UPat(Ops.STACK, name="off")).store(UPat(Ops.STACK, name="val")), fold_const_store),
+  (UPat({Ops.BUFFER, Ops.MSTACK}, name="buf").index(UPat(Ops.STACK, name="off")).store(UPat(Ops.STACK, name="val")), fold_const_store),
 ])

 pm_assert_no_afters = PatternMatcher([(UPat(Ops.AFTER, name="a"), lambda a: panic(RuntimeError, f"AFTER left at hcq_link: {a.src[0].op}"))])
@@ -623,6 +624,8 @@ class HCQAllocator(LRUAllocator[HCQDeviceType], Generic[HCQDeviceType]):
    if not hasattr(self, '_do_map'): raise NotImplementedError("map failed: no method implemented")
    return self._do_map(buf)

+  def _do_unmap(self, mb): self.dev.iface.free(mb)
+
  @suppress_finalizing
  def _free(self, buf:HCQ2Buffer, options:BufferSpec|None=None):
    if options is not None and options.external_ptr is not None: return
@@ -631,6 +634,6 @@ class HCQAllocator(LRUAllocator[HCQDeviceType], Generic[HCQDeviceType]):

  def _unmap(self, mb):
    self.dev.synchronize()
-    self.dev.iface.free(mb)
+    self._do_unmap(mb)

  def _offset(self, buf, size:int, offset:int) -> HCQ2Buffer: return buf.offset(offset=offset, size=size)
@@ -1,7 +1,8 @@
 from __future__ import annotations
-import os, mmap, array, functools, ctypes, select, contextlib, dataclasses, sys, itertools, struct, socket, subprocess, time, enum, atexit
+import os, mmap, array, functools, ctypes, ctypes.util, select, contextlib, dataclasses, sys, itertools, struct, socket
+import subprocess, time, enum, atexit
 from tinygrad.helpers import round_up, getenv, OSX, temp, ceildiv, unwrap, fetch, system, _ensure_downloads_dir, DEBUG, flatten, pluralize
-from tinygrad.runtime.autogen import libc, pci, vfio, iokit, corefoundation
+from tinygrad.runtime.autogen import libc, pci, vfio
 from tinygrad.runtime.support.hcq import FileIOInterface, MMIOInterface, HCQBuffer, hcq_filter_visible_devices
 from tinygrad.runtime.support.memory import VirtMapping, AddrSpace, BumpAllocator
 from tinygrad.runtime.support.usb import USB3, CustomASM24Controller, USBMMIOInterface
@@ -55,6 +56,7 @@ class _System:
  def pci_scan_bus(self, vendor:int, devices:tuple[tuple[int, tuple[int, ...]], ...], base_class:int|None=None) -> list[str]:
    all_devs = []
    if OSX:
+      from tinygrad.runtime.autogen import iokit, corefoundation
      def read_prop(svc, key) -> int:
        cfkey = corefoundation.CFStringCreateWithCString(None, key.encode(), corefoundation.kCFStringEncodingUTF8)
        cfdata = ctypes.cast(iokit.IORegistryEntryCreateCFProperty(svc, ctypes.cast(cfkey, iokit.CFStringRef), None, 0), corefoundation.CFDataRef)
@@ -23,7 +23,7 @@ class IndexingContext:
    return UOp.range(s, next(self.range_idx), axistype) if resolve(s!=1) else UOp.const(0)


-ALWAYS_CONTIGUOUS: set[Ops] = {Ops.CONTIGUOUS, Ops.AFTER, Ops.BUFFER, Ops.SLICE,
+ALWAYS_CONTIGUOUS: set[Ops] = {Ops.CONTIGUOUS, Ops.AFTER, Ops.BUFFER,
                      Ops.CONST, Ops.BIND, Ops.MSELECT, Ops.MSTACK, Ops.PARAM,
                      Ops.LOAD, Ops.CALL, Ops.FUNCTION}

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

 def realize_store_after_src(ctx:IndexingContext, dest:UOp, src:UOp):
-  # don't realize SLICE when it's the direct source of STORE+AFTER — the target buffer is the output
-  if src.op is Ops.SLICE and src in ctx.realize_map \
-     and not dest.op_in_backward_slice_with_self(Ops.SHRINK, Ops.PERMUTE, Ops.FLIP, Ops.PAD):
-    del ctx.realize_map[src]
  # you don't usually have to do this for assign unless there's a WAR hazard like TestAssign.test_assign_double_diamond_reduce
  if dest.base in src.backward_slice_with_self: ctx.realize_map[src] = None

@@ -74,7 +70,7 @@ def broadcast_rngs(x:UOp, src:UOp, rngs:tuple[UOp, ...]) -> tuple[UOp, ...]:
 # TODO: srcs contain (real data srcs, something else, ranges) and the boundary is confusing. see range_start
 def data_srcs(op:Ops, src:tuple[UOp, ...]) -> tuple[UOp, ...]:
  if op in {Ops.PARAM, Ops.BUFFER, Ops.RANGE, Ops.SPECIAL, Ops.BIND}: return ()
-  if op in GroupOp.Movement|{Ops.INDEX, Ops.SLICE, Ops.STAGE, Ops.REDUCE, Ops.AFTER, Ops.END}: return src[:1]
+  if op in GroupOp.Movement|{Ops.INDEX, Ops.STAGE, Ops.REDUCE, Ops.AFTER, Ops.END}: return src[:1]
  return src

 def create_bufferize_and_index_srcs(ctx:IndexingContext, x:UOp) -> list[UOp]:
@@ -84,7 +80,7 @@ def create_bufferize_and_index_srcs(ctx:IndexingContext, x:UOp) -> list[UOp]:
  for i, s in enumerate(x.src):
    new_src = s
    src_rngs = broadcast_rngs(x, s, ctx.range_map[x][0]) if x in ctx.range_map else ()
-    if s.op in {Ops.PARAM, Ops.BUFFER, Ops.SLICE, Ops.MSTACK, Ops.MSELECT, Ops.AFTER}:
+    if s.op in {Ops.PARAM, Ops.BUFFER, Ops.MSTACK, Ops.MSELECT, Ops.AFTER}:
      if x in ctx.range_map and i < data_src_count: new_src = new_src.index(*src_rngs)
    elif s in ctx.realize_map:
      realized_ranges = ctx.realize_map[s]
@@ -52,11 +52,9 @@ def memory_plan_rewrite(linear:UOp, held_bufs:set[UOp]|None=None) -> UOp:
    peaks[_key(buf)] = (max(peaks[_key(buf)][0], offsets[buf] + buf.max_numel() * buf.dtype.itemsize), peaks[_key(buf)][1])
  arena_sizes = {key: round_up(peak, block_size) for key, (peak, _) in peaks.items()}

-  # build replace_map: each buffer becomes a SLICE into a shared per-device-lane arena
+  # build replace_map: each buffer becomes a SHRINK/BITCAST into a shared per-device-lane arena
  arenas = {key: UOp.new_buffer(key[0], sz, dtypes.int8) for key, sz in arena_sizes.items()}
-  replace_map:dict[UOp, UOp] = {}
-  for buf_uop, offset in offsets.items():
-    replace_map[buf_uop] = UOp(Ops.SLICE, buf_uop.dtype, (arenas[_key(buf_uop)], UOp.const(offset)), buf_uop.max_numel())
+  replace_map = {buf_uop:arenas[_key(buf_uop)][offset:offset+buf_uop.nbytes()].bitcast(buf_uop.dtype) for buf_uop, offset in offsets.items()}

  if DEBUG >= 1 and (omem:=sum(nbytes.values()) / 1e6) != (nmem:=sum(arena_sizes.values()) / 1e6):
    print(f"memory reduced from {omem:.2f} MB -> {nmem:.2f} MB, {len(first_appearance)} -> {len(arenas)} bufs")
@@ -6,11 +6,11 @@ from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, K
 from tinygrad.uop.ops import graph_rewrite, sint, AxisType, BottomUpGate, rewrite_group, identity_element
 from tinygrad.uop.symbolic import symbolic, pm_fold_cast_const
 from tinygrad.uop.movement import mop_cleanup
-from tinygrad.helpers import prod, getenv, dedup, all_int, DEBUG, SPLIT_REDUCEOP, VIZ, MAX_KERNEL_BUFFERS, SPEC
+from tinygrad.helpers import prod, getenv, dedup, all_int, DEBUG, SPLIT_REDUCEOP, DEBUG_RANGEIFY, VIZ, MAX_KERNEL_BUFFERS, SPEC
 from tinygrad.helpers import PCONTIG, FLOAT16, OPENPILOT_HACKS, argsort, partition, get_single_element
 from tinygrad.codegen.simplify import pm_flatten_range, pm_reduce_simplify
 from tinygrad.codegen.opt import Opt
-from tinygrad.schedule.indexing import BufferizeOpts, IndexingContext, apply_movement_op
+from tinygrad.schedule.indexing import run_rangeify, BufferizeOpts, IndexingContext, apply_movement_op
 from tinygrad.schedule.multi import multi_pm
 from tinygrad.schedule.allreduce import create_allreduce_function

@@ -39,16 +39,7 @@ pm_fold_moved_after = PatternMatcher([
 def _mop_index(r:UOp, idx:UOp):
  idxs = idx.src[1:]
  if len(idxs) == len(r.shape):
-    ret = r.src[0].index(*apply_movement_op(r.op, r.src[0].shape, r.marg, idxs), dtype=idx.dtype, arg=idx.arg)
-    if r.op is Ops.PAD:
-      # insert 0 for PAD with where
-      # TODO: does this need simplify
-      a = UOp.const(True)
-      for s in ret.src[1:]:
-        if s.op is Ops.WHERE and s.src[2].op is Ops.CONST and s.src[2].arg == Invalid:
-          a = a & s.src[0]
-      ret = a.where(ret, ret.const_like(0))
-    return ret
+    return r.src[0].index(*apply_movement_op(r.op, r.src[0].shape, r.marg, idxs), dtype=idx.dtype, arg=idx.arg)
  if r.op is Ops.RESHAPE:
    src_prefix = len(r.src[0].shape) - len(r.shape[len(idxs):])
    if src_prefix >= 0 and r.src[0].shape[src_prefix:] == r.shape[len(idxs):]:
@@ -576,101 +567,9 @@ def convert_copy_to_store(ctx, copy:UOp, existing_buf:UOp|None=None):
  # reshape back to input
  return buf.after(buf.store(input_src)).reshape(copy.shape)

-def convert_contig_to_store(ctx, copy:UOp):
-  input_src = copy.src[0]
-  # create the output buffer
-  buf = UOp(Ops.BUFFER, src=(shape_to_shape_arg(input_src.max_shape),), arg=ParamArg(next(ctx), copy.dtype, device=copy.device))
-  # reshape back to input
-  view = buf.shrink_to(input_src.shape)
-  return view.after(view.store(input_src))
-
 pm_copy_to_store = PatternMatcher([
  (UPat(name="existing_buf").store(UPat(Ops.COPY, name="copy")), convert_copy_to_store),
  (UPat(Ops.COPY, name="copy"), convert_copy_to_store),
-  (UPat(Ops.CONTIGUOUS, name="copy"), convert_contig_to_store),
-])
-
-# **** simple rangeify ****
-
-from tinygrad.helpers import all_same
-from tinygrad.uop.ops import _broadcast_shape
-
-def expand_broadcast(x:UOp):
-  shapes = [u._shape for u in x.src]
-  if any(s is None for s in shapes) or all_same(shapes): return None
-  shape = _broadcast_shape(*shapes)
-  return x.replace(src=tuple([u.expand(shape) for u in x.src]))
-
-pm_expand_broadcast = PatternMatcher([
-  # expand broadcasts first
-  (UPat(GroupOp.Binary|GroupOp.Ternary|{Ops.STORE}, name="x"), expand_broadcast),
-])
-
-def expand_coeff(sink:UOp) -> dict[UOp,int]:
-  coeff: dict[UOp,int] = {sink: 1}
-  contig: dict[UOp,int] = {}
-  for u in reversed(list(sink.toposort())):
-    c = 1 if u.op is Ops.STORE else coeff.get(u, 0)
-    # symbolic coeffs mark on vmax, an extra CONTIGUOUS is always safe
-    if (c > 1 if isinstance(c, int) else c.vmax > 1) and u.op in (GroupOp.Elementwise | {Ops.REDUCE}) and u.device is not None:
-      contig[u] = c
-      c = 1
-    coeff[u] = c
-    mult = prod(u.shape) // prod(u.src[0].shape) if u.op is Ops.EXPAND else 1
-    for s in u.src: coeff[s] = coeff.get(s, 0) + c * (mult if s is u.src[0] else 1)
-  return contig
-
-def rangeify_on_reduce(ctx, inp:UOp, red:UOp, idx:UOp|None=None):
-  if red.arg[1] == 0: return None
-  if idx is None and len(red.shape) > 0: return None
-  # TODO: is AxisType.REDUCE a real thing?
-  rngs = [UOp.range(s, next(ctx), AxisType.REDUCE) for s in inp.shape[:red.arg[1]]]
-  return inp.index(*rngs, *(idx.src[1:] if idx is not None else ())).reduce(*rngs, arg=(red.arg[0], 0))
-
-def rangeify_on_store(ctx, x:UOp):
-  if x.shape == (): return None
-  rngs = [UOp.range(s, next(ctx)) for s in x.shape]
-  return x.src[0].index(*rngs).store(x.src[1].index(*rngs)).end(*rngs)
-
-def rangeify_on_stage(ctx, x:UOp):
-  if x.src[0].shape == (): return None
-  # size 1 dims don't get ranges, they are reshaped out and back in
-  if all_int(x.shape) and 0 < len(sq := tuple(s for s in x.shape if s != 1)) < len(x.shape):
-    return rangeify_on_stage(ctx, x.src[0].reshape(sq).bufferize(arg=x.arg)).reshape(x.shape)
-  rngs = [UOp.range(s, next(ctx)) for s in x.shape]
-  return x.replace(src=(x.src[0].index(*rngs), *rngs))
-
-def index_on_stack(stack:UOp, idx:UOp):
-  srcs = [s.index(*idx.src[2:]) for s in stack.src]
-  r0 = idx.src[1]
-  ret = srcs[-1]
-  for k in range(len(srcs)-2, -1, -1): ret = r0.eq(k).where(srcs[k], ret)
-  return ret
-
-pm_simple_rangeify = PatternMatcher([
-  # INDEX without src is nothing
-  (UPat(Ops.INDEX, src=(UPat.var('x'),)), lambda x: x),
-  # STAGE on shape () is nothing
-  (UPat(Ops.STAGE, src=(UPat.var('x'),)), lambda x: x if x.shape == () else None),
-  # if INDEX is on STAGE with the same ranges, remove the pair
-  (UPat(Ops.STAGE, allow_any_len=True, name="s").index(allow_any_len=True, name="i"),
-   lambda s,i: s.src[0] if s.src[1:] == i.src[1:] else None),
-  # reshape of a single element shaped value to scalar is an index
-  (UPat(Ops.RESHAPE, name="x"), lambda x: x.src[0].index(0) if x.marg == () and x.src[0].shape == (1,) else None),
-  # handle movement ops on INDEX
-  (UPat(GroupOp.Movement, name="r").index(name="idx", allow_any_len=True), _mop_index),
-  (UPat(Ops.STACK, name="stack").index(name="idx", allow_any_len=True), index_on_stack),
-  # pass index through elementwise
-  (UPat(GroupOp.Elementwise, name="b").index(name="idx", allow_any_len=True),
-   lambda b,idx: b.replace(src=tuple(s.index(*idx.src[1:]) for s in b.src))),
-])
-
-pm_range_creation = PatternMatcher([
-  # reduce/store are what creates ranges
-  (UPat(Ops.REDUCE, src=(UPat.var('inp'),), name="red").index(name="idx", allow_any_len=True), rangeify_on_reduce),
-  (UPat(Ops.REDUCE, src=(UPat.var('inp'),), name="red"), rangeify_on_reduce),
-  (UPat(Ops.STORE, name="x"), rangeify_on_store),
-  (UPat(Ops.STAGE, name="x"), rangeify_on_stage),
 ])

@rewrite_group(new_ctx=False)
@@ -679,61 +578,15 @@ def get_kernel_graph(sink:UOp) -> UOp:
  if OPENPILOT_HACKS: tsink = graph_rewrite(tsink, pm_fold_moved_after, ctx={}, name="fold moved afters")
  tsink = graph_rewrite(tsink, pm_mops+earliest_rewrites, bottom_up=True, name="earliest rewrites")

-  # convert movement ops to ranges
-  #tsink, rctx = run_rangeify(tsink, bool(DEBUG_RANGEIFY))
-  tsink = graph_rewrite(tsink, pm_expand_broadcast, bottom_up=True, name="expand broadcast")
-
-  # mark ops that would be recomputed (expand coeff > 1) as CONTIGUOUS, like the realize map in run_rangeify
-  contig = expand_coeff(tsink)
-  subs: dict[UOp, UOp] = {}
-  for u in tsink.toposort():
-    u2 = u.replace(src=tuple(subs.get(s, s) for s in u.src))
-    subs[u] = u2.alu(Ops.STAGE, arg=BufferizeOpts(u2.device)) if u in contig else u2
-  tsink = subs[tsink]
-
-  # add buffers on copy
  tsink = graph_rewrite(tsink, pm_copy_to_store, ctx=itertools.count(0), bottom_up=True, name="convert copy to store")

-  # simple rangeify
-  tsink = graph_rewrite(tsink, pm_range_creation+pm_simple_rangeify, ctx=itertools.count(0), bottom_up=True, name="simple rangeify")
+  # convert movement ops to ranges
+  tsink, rctx = run_rangeify(tsink, bool(DEBUG_RANGEIFY))

-  # for each index on a stage without children, try to merge the stage into the consumer kernel
-  while 1:
-    indexes: dict[UOp, list[UOp]] = {}
-    consumers: dict[UOp, list[UOp]] = {}
-    for u in tsink.toposort():
-      if u.op is Ops.INDEX and u.src[0].op is Ops.STAGE:
-        indexes.setdefault(u.src[0], []).append(u)
-      for s in u.src: consumers.setdefault(s, []).append(u)
-    # the ranges wrapping u: REDUCE ranges on the path up, plus the enclosing END/STAGE nest
-    def nest_ranges(u:UOp) -> set[UOp]:
-      ret: set[UOp] = set()
-      stack, seen = [u], set()
-      while len(stack):
-        if (x := stack.pop()) in seen: continue
-        seen.add(x)
-        if x.op is Ops.REDUCE: ret.update(*[er.ranges for er in x.ended_ranges])
-        elif x.op in {Ops.END, Ops.STAGE}:
-          ret.update(*[er.ranges for er in x.ended_ranges])
-          continue
-        stack.extend(consumers.get(x, []))
-      return ret
-    subs = {}
-    for k,v in indexes.items():
-      # don't move REDUCE ranges up (real?)
-      if len(v) != 1 or not all(all([r.arg[-1] == AxisType.WEAK for r in s.ranges]) for s in v[0].src[1:]): continue
-      # merging must not add iteration multiplicity around range-bound computation in the stage body:
-      # every range of the enclosing kernel nest must be used by the index, unless the body has no inner loops
-      if not nest_ranges(v[0]) <= set().union(*[s.ranges for s in v[0].src[1:]]) and \
-         any(x.op is Ops.REDUCE for x in k.src[0].toposort(gate=lambda x: x.op is not Ops.STAGE)): continue
-      for old_r, new_r in zip(k.src[1:], v[0].src[1:]):
-        subs[old_r] = new_r
-    if not len(subs): break
-    tsink = tsink.substitute(subs)
-    tsink = graph_rewrite(tsink, pm_simple_rangeify, bottom_up=True, name=f"merge kernels ({len(subs)})")
-
-  tsink = graph_rewrite(tsink, symbolic+pm_reduce_simplify+pm_const_buffer_folding+pm_remove_bufferize+pm_no_indexing_calls, name="symbolic+reduce_collapse+debuf")
-  #tsink = graph_rewrite(tsink, pm_limit_bufs, ctx=rctx, name="limit buffers")
+  tsink = graph_rewrite(tsink,
+                        symbolic+pm_fold_cast_const+pm_reduce_simplify+pm_const_buffer_folding+pm_remove_bufferize+pm_no_indexing_calls,
+                        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")

@@ -23,6 +23,7 @@ class AllocCtx:
  bases: set[UOp] = field(default_factory=set)
  assigns: list[UOp] = field(default_factory=list)
  replacements: list[UOp] = field(default_factory=list)
+  views: set[UOp] = field(default_factory=set)

 def tag_uop(ctx:AllocCtx, x:UOp):
  if x.tag is not None: return None
@@ -63,40 +64,37 @@ def replace_contig_with_store_after(u:UOp):
 def replace_store_after_with_contig(u:UOp, src:UOp):
  assigned_to = u
  while assigned_to.op in {Ops.BITCAST, Ops.AFTER, Ops.UNSHARD}: assigned_to = assigned_to.src[0].base
-  if assigned_to.op not in {Ops.BUFFER, Ops.SLICE}: return src.contiguous(tag=u.tag)
+  if assigned_to.op is not Ops.BUFFER: return src.contiguous(tag=u.tag)

 def _make_buffer_view(src:UOp) -> UOp|None:
-  """If movement ops on src collapse to a contiguous range, return SLICE. Otherwise None."""
-  if (offset := src.contiguous_view_offset()) is None: return None
-  buf = src.base
-  if buf.op is Ops.SLICE:
-    byte_offset = buf.src[1].val * buf.src[0].dtype.itemsize + offset * src.dtype.itemsize
-    buf = buf.src[0]
-    if byte_offset % buf.dtype.itemsize != 0: return None
-    offset = byte_offset // buf.dtype.itemsize
-  return UOp(Ops.SLICE, src.dtype, (buf, UOp.const(offset)), src.numel())
+  if (cv := src.contiguous_view()) is None: return None
+  (buf, offset), size = cv, src.max_numel() * src.element_size() // cv[0].element_size()
+  if buf.op is not Ops.BUFFER: return None
+  # NB: make offset a UOp.variable here to do the offset computation in the kernels
+  return buf[offset:offset+size].bitcast(src.dtype)

-def contiguous_mops_to_view(c:UOp, src:UOp):
-  """MOPS(BUFFER) → SLICE when movement ops collapse to a contiguous range."""
+def contiguous_mops_to_view(ctx:AllocCtx, c:UOp, src:UOp):
+  """MOPS(BUFFER) → SHRINK when movement ops collapse to a contiguous range."""
  buf = src.base
-  if buf.op not in {Ops.BUFFER, Ops.SLICE, Ops.UNSHARD}: return None
-  if src.op is Ops.RESHAPE and src.src[0].op in {Ops.BUFFER, Ops.SLICE} and c.op is not Ops.BITCAST: return None
-  if c.op is not Ops.BITCAST and src.op is Ops.BUFFER: return None
+  while buf.op is Ops.BITCAST: buf = buf.src[0].base
+  if buf.op not in {Ops.BUFFER, Ops.UNSHARD}: return None

  # no symbolic shape
  if not all_int(c.shape): return None

  if buf.op is not Ops.UNSHARD and (view := _make_buffer_view(src)) is not None:
-    view = (view.replace(dtype=c.dtype, arg=c.numel()) if c.op is Ops.BITCAST else view).reshape(c.shape)
-    return c.replace(src=(view,)) if c.op is Ops.COPY else view
+    ctx.views.add(view)
+    view = view.reshape(c.shape)
+    return c.replace(src=(view,)+c.src[1:]) if c.op in {Ops.COPY, Ops.STORE} else view

-  # for UNSHARD tensors, use multi_pm to resolve per-shard movement ops, then create SLICE on the resolved result
+  # for UNSHARD tensors, use multi_pm to resolve per-shard movement ops, then create SHRINK on the resolved result
  if not isinstance(c.device, str):
    from tinygrad.schedule.multi import multi_pm
    resolved = graph_rewrite(src, multi_pm, name="multi_buffer_view")
    if resolved.op is not Ops.UNSHARD: return None
    if (view := _make_buffer_view(resolved.src[0])) is None: return None
-    return view.reshape(resolved.src[0].shape).unshard(resolved.arg, resolved.src[1:]).contiguous(tag=c.tag)
+    ctx.views.add(view)
+    return view.reshape(resolved.src[0].shape).unshard(resolved.arg, resolved.src[1:])

  return None

@@ -151,8 +149,9 @@ pm_early_transform_tensor_graph = PatternMatcher([
  # resolve TUPLE+GETTUPLE (for precompiled calls)
  (UPat(Ops.GETTUPLE, src=(UPat(Ops.TUPLE, name="t"),), name="g"), lambda g,t: t.src[g.arg]),

-  # fold MOPS+BITCAST over BUFFER/SLICE into SLICE when movement ops collapse to contiguous range
-  (UPat((Ops.BITCAST, Ops.COPY, Ops.CONTIGUOUS), src=(UPat(GroupOp.Movement|{Ops.BUFFER}, name="src"),), name="c"), contiguous_mops_to_view),
+  # fold MOPS+BITCAST over BUFFER into SHRINK when movement ops collapse to contiguous range
+  (UPat((Ops.COPY, Ops.CONTIGUOUS), src=(UPat(GroupOp.Movement|{Ops.BITCAST}, name="src"),), name="c"), contiguous_mops_to_view),
+  (UPat(Ops.STORE, src=(UPat(Ops.BITCAST, name="src"), UPat()), name="c", allow_any_len=True), contiguous_mops_to_view),

  # 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, name="copy"), lambda x,copy:
@@ -201,6 +200,8 @@ def replace_input_buffer(ctx:AllocCtx, b:UOp):
  return UOp.param(len(ctx.replacements)-1, b.dtype, b.shape, b.device,
                   addrspace=b.addrspace if b.addrspace is not None else AddrSpace.GLOBAL)

+def replace_input_view(ctx:AllocCtx, b:UOp): return replace_input_buffer(ctx, b) if b in ctx.views else None
+
 pm_finalize_call = PatternMatcher([
  (UPat(Ops.AFTER, name="x"), finalize_after),
  (UPat(Ops.COPY, name="x"), lambda ctx,x: ctx.assigns.append(x) if isinstance(x.device, str) and x.device.startswith(("DISK", "TINYFS")) else None),
@@ -210,8 +211,9 @@ pm_replace_buf = PatternMatcher([
  # replace BUFFER with PARAM for cache key normalization
  (UPat(Ops.BUFFER, src=(UPat(),), name="b"), lambda ctx,b:
   replace_input_buffer(ctx, b) if isinstance(b.arg, ParamArg) and b.addrspace is AddrSpace.GLOBAL else None),
-  # replace SLICE with PARAM. this rewrite is bottom up so BUFFERs we don't need won't be in the input
-  (UPat(Ops.SLICE, src=(UPat(Ops.BUFFER), UPat(Ops.CONST, dtype=dtypes.weakint)), name="b"), replace_input_buffer),
+  # replace SHRINK with PARAM
+  (UPat(Ops.SHRINK, src=(UPat(Ops.BUFFER),), name="b", allow_any_len=True), replace_input_view),
+  (UPat(Ops.BITCAST, src=(UPat.any(UPat(Ops.SHRINK, src=(UPat(Ops.BUFFER),), allow_any_len=True), UPat(Ops.BUFFER)),), name="b"), replace_input_view),
  # strip value from BIND for cache key normalization, so different values hit same cache
  (UPat(Ops.BIND, src=(UPat(Ops.PARAM), UPat(Ops.CONST)), name="b"), replace_input_buffer),
 ])
@@ -229,7 +231,7 @@ def transform_to_call(big_sink:UOp) -> tuple[UOp, dict[UOp, UOp]]:
  big_sink = graph_rewrite(big_sink, add_tags, ctx=ctx, bottom_up=True, name="number the uops")

  # here we can break the tensor graph. this is the only place you need to maintain numbered tags
-  big_sink = graph_rewrite(big_sink, pm_early_transform_tensor_graph, name="early transform tensor graph")
+  big_sink = graph_rewrite(big_sink, pm_early_transform_tensor_graph, ctx=ctx, name="early transform tensor graph")

  # here we construct the final buffer_map: as-built nodes -> their final storage. values are never keys
  graph_rewrite(big_sink, pm_finalize_call, ctx=ctx, name="finalize call")
@@ -93,7 +93,7 @@ class Ops(FastEnum):
  CONTIGUOUS = auto(); CONTIGUOUS_BACKWARD = auto(); DETACH = auto()

  # buffer ops
-  STAGE = auto(); COPY = auto(); SLICE = auto(); MSELECT = auto(); MSTACK = auto(); CUSTOM_FUNCTION = auto()
+  STAGE = auto(); COPY = auto(); MSELECT = auto(); MSTACK = auto(); CUSTOM_FUNCTION = auto()

  # the core 6 movement ops! these only exist in the tensor graph
  RESHAPE = auto(); PERMUTE = auto(); EXPAND = auto(); PAD = auto(); FLIP = auto()
@@ -45,8 +45,7 @@ axis_colors = {AxisType.DEVICE: "green", AxisType.GLOBAL: "blue", AxisType.THREA
 axis_to_pos = {AxisType.DEVICE: -2, AxisType.WEAK: -1, AxisType.LOOP: -1, AxisType.THREAD: 0, AxisType.GLOBAL: 0, AxisType.WARP: 1,
               AxisType.LOCAL: 2, AxisType.UPCAST: 3, AxisType.GROUP_REDUCE: 2, AxisType.REDUCE: 4, AxisType.UNROLL: 5}

-range_start = {Ops.STAGE: 1, Ops.REDUCE: 1, Ops.WMMA: 3, Ops.END: 1, Ops.CALL: 1, Ops.FUNCTION: 1,
-               Ops.SLICE: 2, Ops.LINEAR: 0}
+range_start = {Ops.STAGE: 1, Ops.REDUCE: 1, Ops.WMMA: 3, Ops.END: 1, Ops.CALL: 1, Ops.FUNCTION: 1, Ops.LINEAR: 0}

 # https://en.wikipedia.org/wiki/Identity_element
 def identity_element(op:Ops, dt:DType) -> PyConst: return dt.const({Ops.ADD:0, Ops.MUL:1, Ops.MAX:dt.min}[op])
@@ -171,9 +170,6 @@ def dtype_from_uop(op:Ops, src:tuple[UOp,...], arg:Any) -> DType|None:
      return arg.dtype
    case Ops.BINARY:
      return dtypes.uint8
-    case Ops.SLICE:
-      # TODO: slice just shouldn't exist
-      return None
    case Ops.CAST | Ops.BITCAST:
      assert isinstance(arg, DType), f"CAST/BITCAST arg must be DType, got {arg}"
      return arg
@@ -221,7 +217,7 @@ class UOpMetaClass(type):
    return created

 # some uops map to other stuff
-buffers:weakref.WeakKeyDictionary[UOp, Buffer|MultiBuffer] = weakref.WeakKeyDictionary() # this maps BUFFER/SLICE uops to their device Buffers
+buffers:weakref.WeakKeyDictionary[UOp, Buffer|MultiBuffer] = weakref.WeakKeyDictionary() # this maps BUFFER/view uops to their device Buffers
 all_metadata:weakref.WeakKeyDictionary[UOp, tuple[Metadata, ...]] = weakref.WeakKeyDictionary() # TODO: should this be here?

 # recursive_property replaces functools.cached_property in recursive UOp functions to prevent RecursionError
@@ -386,10 +382,6 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
      case Ops.BUFFER:
        if len(self.src): return self.src[0].as_shape
        return ()
-      case Ops.SLICE:
-        # HACK: SLICE is used inside kernels, so we set the shape to () if it's on an INDEX
-        if self.src[0].op is Ops.INDEX: return ()
-        return (self.arg,)
      case Ops.CUSTOM | Ops.CUSTOMI:
        if self.dtype is dtypes.void: return None
        input_shapes = [x._shape for x in self.src if x._shape is not None]
@@ -568,9 +560,9 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
    in_tuple = self.src[0] if self.op is Ops.FUNCTION else self
    assert in_tuple.op is Ops.TUPLE, f"gettuple requires FUNCTION or TUPLE source, got {self.op}"
    return UOp(Ops.GETTUPLE, src=(self,), arg=idx)
-  def group(*srcs:UOp|None):  # pylint: disable=no-self-argument
+  def group(*srcs:UOp|None, **kwargs):  # pylint: disable=no-self-argument
    if len(srcs) == 1 and isinstance(srcs[0], UOp): return srcs[0]
-    return UOp(Ops.GROUP, src=tuple([x for x in srcs if x is not None]))
+    return UOp(Ops.GROUP, src=tuple([x for x in srcs if x is not None]), **kwargs)
  def index(self, *srcs:UOp|int|None, **kwargs):
    new_srcs: list[UOp] = [UOp.const(x) if isinstance(x, int) else x for x in srcs if x is not None]
    if len(new_srcs) == 1 and new_srcs[0].op is Ops.CONST and self.op is Ops.STACK: return self.src[new_srcs[0].val]
@@ -822,7 +814,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
  unique_num = itertools.count(0)

  def getaddr(self, device=None) -> UOp:
-    if self.without_after.op not in {Ops.BUFFER, Ops.SLICE, Ops.BINARY, Ops.MSTACK, Ops.MSELECT, Ops.PARAM}: return self
+    if self.without_after.op not in {Ops.BUFFER, Ops.SHRINK, Ops.BITCAST, Ops.BINARY, Ops.MSTACK, Ops.MSELECT, Ops.PARAM}: return self
    return UOp(Ops.GETADDR, src=(self,), arg=device or to_tuple(self.device)[0])
  @staticmethod
  def new_buffer(device:str|tuple[str, ...], size:int, dtype:DType, num=None):
@@ -901,8 +893,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
    while len(s.src) and s.op not in {Ops.BUFFER, Ops.PARAM, Ops.STAGE, Ops.MSTACK}: s = s.src[0]
    return s

-  def contiguous_view_offset(self) -> int|None:
-    """If movement ops on a BUFFER collapse to a contiguous range, return `offset` in elements. Otherwise None."""
+  def contiguous_view(self) -> tuple[UOp, int]|None:
    from tinygrad.schedule.rangeify import pm_mops
    from tinygrad.uop.symbolic import symbolic

@@ -915,14 +906,17 @@ class UOp(RandMixin, metaclass=UOpMetaClass):

    idx = self.flatten().index(UOp.range(self.numel(), 0))
    out = graph_rewrite(idx, pm_mops+symbolic+pm_contiguous_view_offset, ctx=self, name="contiguous_view_offset")
-    return out.val if out.op is Ops.CONST and isinstance(out.val, int) else None
+    if out.op is not Ops.INDEX or not (b:=out.src[0]).tag or (c:=out.src[1]).op is not Ops.CONST or not isinstance(c.val, int): return None
+    return b.rtag(None), c.val
+
+  def contiguous_view_offset(self) -> int|None: return None if (view := self.contiguous_view()) is None else view[1]

  def has_buffer_identity(self, after_ok=False):
    """Check if this UOp has a concrete buffer identity in the graph (RESHAPE/UNSHARD -> BUFFER chain)."""
    # TODO: this is confusing because UOp.variable('v', 0, 1, dtypes.weakfloat) is True for jit to work, but it doesn't have a buffer
    if self.op in {Ops.RESHAPE, Ops.UNSHARD, Ops.MSELECT}: return self.src[0].has_buffer_identity(after_ok)
    if after_ok and self.op == Ops.AFTER: return self.src[0].has_buffer_identity(after_ok)
-    return self.op in {Ops.BUFFER, Ops.SLICE, Ops.PARAM}
+    return self.op in {Ops.BUFFER, Ops.PARAM}

  def _base_buffer_is_realized(self) -> bool:
    """Walk through AFTER chain to find if the underlying buffer is realized (has allocated memory)."""
@@ -932,30 +926,19 @@ class UOp(RandMixin, metaclass=UOpMetaClass):

  @property
  def buffer(self) -> Buffer|MultiBuffer:
-    if self.op in {Ops.CONTIGUOUS, Ops.RESHAPE, Ops.UNSHARD, Ops.DETACH, Ops.AFTER}: return self.src[0].buffer
+    if self.op in {Ops.CONTIGUOUS, Ops.CONTIGUOUS_BACKWARD, Ops.RESHAPE, Ops.UNSHARD, Ops.DETACH, Ops.AFTER}: return self.src[0].buffer
    # this buffer can process disk tensors and simple movement ops
-    if self is not self.base:
-      buf = self.base.buffer
-      assert isinstance(buf, Buffer), "must be a Buffer for movement ops"
-      offset = self.contiguous_view_offset()
-      if offset is None: raise RuntimeError(f"non-contiguous view is not supported for {buf.device} buffer")
-      return buf.view(prod(self.max_shape), self.dtype, offset*self.dtype.itemsize)
-    if self.op is Ops.BITCAST:
-      buf = self.src[0].buffer
-      assert isinstance(buf, Buffer), "must be a Buffer for BITCAST"
-      return buf.view(prod(self.max_shape), self.dtype, 0)
-    if self.op is Ops.SLICE:
+    if self is not self.base or self.op is Ops.BITCAST:
      if (cret:=buffers.get(self)) is not None: return cret
-      buf = self.src[0].buffer
-      offset = self.src[1].val
+      if (cv := self.contiguous_view()) is None: raise RuntimeError(f"non-contiguous view is not supported for {self.device} buffer")
+      buf, offset = (b:=cv[0]).base.buffer, cv[1]
      if isinstance(buf, MultiBuffer):
        mbuf = MultiBuffer.__new__(MultiBuffer)
-        mbuf.bufs = [b.view(self.arg, self.dtype, offset * self.src[0].dtype.itemsize) for b in buf.bufs]
+        mbuf.bufs = [x.view(prod(self.max_shape), self.dtype, offset*b.dtype.itemsize) for x in buf.bufs]
        buffers[self] = mbuf
        return mbuf
-      assert isinstance(buf, Buffer), "must be a Buffer for SLICE"
-      buffers[self] = bv = buf.view(self.arg, self.dtype, offset * self.src[0].dtype.itemsize)
-      return bv
+      buffers[self] = buf.view(prod(self.max_shape), self.dtype, offset*b.dtype.itemsize)
+      return buffers[self]
    if self.op is Ops.MSELECT:
      ret = self.src[0].buffer
      assert isinstance(ret, MultiBuffer)
@@ -1181,7 +1164,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
  def custom_function(name:str, *src:UOp) -> UOp: return UOp(Ops.CUSTOM_FUNCTION, src=src, arg=name)

  # opaque bodies stay as Ops.CALL; value-producing bodies become Ops.FUNCTION (wrapped in TUPLE)
-  _OPAQUE_CALL_BODIES = {Ops.SINK, Ops.PROGRAM, Ops.LINEAR, Ops.COPY, Ops.SLICE, Ops.CUSTOM_FUNCTION}
+  _OPAQUE_CALL_BODIES = {Ops.SINK, Ops.PROGRAM, Ops.LINEAR, Ops.COPY, Ops.CUSTOM_FUNCTION}
  def call(self, *srcs:UOp, ret_dtype:DType|None=None, grad_fxn:Callable|None=None,
           name:str|None=None, precompile:bool=False, precompile_backward:bool=False, aux:Any=None) -> UOp:
    if ret_dtype is not None: return UOp(Ops.CALL, ret_dtype, src=(self,)+srcs)
@@ -1775,10 +1758,14 @@ pm_unbind = PatternMatcher([(UPat(Ops.BIND, name="x"), do_unbind)])

 # ctx is source UOp for which we are finding a contiguous view for. used in contiguous_view_offset
 pm_contiguous_view_offset = PatternMatcher([
-  (UPat(Ops.INDEX, src=(UPat(),)), lambda: UOp.const(0)),
-  (UPat(Ops.INDEX, src=(UPat(), UPat(Ops.RANGE))), lambda: UOp.const(0)),
-  (UPat(Ops.INDEX, src=(UPat(), UPat(Ops.RANGE)+UPat.cvar('c'))), lambda c: c),
-  (UPat(Ops.INDEX, src=(UPat(), UPat.cvar('c'))), lambda ctx, c: c if resolve(ctx.numel() == 1, False) else None),
+  # normalize to 1d bitcasts
+  (UPat(Ops.BITCAST, name="b"), lambda b: b.src[0].flatten().bitcast(b.dtype).reshape(b.shape) if len(b.shape) != 1 else None),
+  (UPat(Ops.BITCAST, name="b").index(UPat.cvar("c")), lambda ctx, b, c:
+   b.src[0].flatten().index(UOp.range(ctx.numel() * (osz:=b.element_size())//(isz:=b.src[0].element_size()), 0) + (c * osz//isz)) if b.tag else None),
+  (UPat(Ops.INDEX, src=(UPat.var("b"),)), lambda b: b.rtag().index(0)),
+  (UPat(Ops.INDEX, src=(UPat.var("b"), UPat(Ops.RANGE))), lambda b: b.rtag().index(0)),
+  (UPat(Ops.INDEX, src=(UPat.var("b"), UPat(Ops.RANGE)+UPat.cvar('c'))), lambda ctx, b, c: b.rtag().index(c)),
+  (UPat(Ops.INDEX, src=(UPat.var("b"), UPat.cvar('c'))), lambda ctx, b, c: b.rtag().index(c) if resolve(ctx.numel() == 1, False) else None),
 ])

 # *** what was symbolic.py ***
@@ -233,13 +233,6 @@ spec_hcq = PatternMatcher([
 spec_full = PatternMatcher([
  (UPat(Ops.REWRITE_ERROR, dtypes.void, name="x"), lambda x: isinstance(x.arg, str)),

-  # SLICE on BUFFER is allowed if BUFFER is
-  (UPat(Ops.SLICE, src=(UPat(GroupOp.Movement.union({Ops.BUFFER, Ops.PARAM, Ops.STAGE, Ops.AFTER})),
-                        UPat(Ops.CONST, dtype=dtypes.weakint)), allow_any_len=True, name="bv"),
-   lambda bv: isinstance(bv.arg, int)),
-
-  (UPat(Ops.CALL, dtypes.void, src=(UPat((Ops.SLICE,)),), allow_any_len=True), lambda: True),
-
  # codegen may end ranges after gpudims has replaced RANGE with SPECIAL.
  (UPat(Ops.END, src=(UPat(), UPat()), allow_any_len=True), lambda: True),

@@ -249,6 +249,8 @@ symbolic = symbolic_simple+commutative+PatternMatcher([
   lambda x: x.const_like(x.vmin) if x.vmin == x.vmax else None),
  (UPat(Ops.RANGE, src=(UPat(Ops.CONST,)), name="x"), lambda x: x.const_like(x.vmin) if x.vmin == x.vmax else None),
  # max folding
+  ((UPat.cvar("a") < UPat.var("b")).where(UPat.var("b"), UPat.cvar("c")), lambda a,b,c: UOp.maximum(a,b) if a.val == c.val else None),
+  ((UPat.var("a") < UPat.cvar("b")).where(UPat.cvar("c"), UPat.var("a")), lambda a,b,c: UOp.maximum(a,b) if b.val == c.val else None),
  (UPat.maximum(UPat.var("x"), UPat.var("y")), lambda x,y: x if x.vmin >= y.vmax else y if x.vmax <= y.vmin else None),
  # TODO: why does this rule break beautiful_mnist?
  #((UPat.var("x")+UPat.var("z")).maximum(UPat.var("y")+UPat.var("z")), lambda x,y,z: x.maximum(y) + z),
@@ -50,7 +50,7 @@ uops_colors = {Ops.LOAD: "#ffc0c0", Ops.STORE: "#87CEEB", Ops.CONST: "#e0e0e0",
               Ops.INDEX: "#CEF9B7", Ops.STACK: "#D8F9E4",
               Ops.WMMA: "#efefc0", Ops.UNSHARD: "#f6ccff", Ops.INS: "#eec4ff",
               **{x:"#D8F9E4" for x in GroupOp.Movement}, **{x:"#ffffc0" for x in GroupOp.ALU}, Ops.THREEFRY:"#ffff80",
-               Ops.SLICE: "#E5EAFF", Ops.BUFFER: "#B0BDFF", Ops.GETADDR: "#9DB1F0", Ops.COPY: "#a040a0", Ops.CUSTOM_FUNCTION: "#bf71b6",
+               Ops.BUFFER: "#B0BDFF", Ops.GETADDR: "#9DB1F0", Ops.COPY: "#a040a0", Ops.CUSTOM_FUNCTION: "#bf71b6",
               Ops.CALL: "#00B7C8", Ops.FUNCTION: "#C07788", Ops.PARAM: "#14686F", Ops.SOURCE: "#c0c0c0", Ops.BINARY: "#404040",
               Ops.LINEAR: "#7DF4FF",
               Ops.ALLREDUCE: "#ff40a0", Ops.MSELECT: "#d040a0", Ops.MSTACK: "#d040a0", Ops.CONTIGUOUS: "#FFC14D",
Author	SHA1	Message	Date
George HotzandGitHub	26f2049d43	Merge branch 'master' into bitcast_spec	2026-08-12 19:24:33 -07:00
geohot	ff0cb28c21	skip slow whisper tests	2026-08-12 13:12:37 -07:00
qazalandGitHub	ed8297a102	kerenl opts test from nan in llama 8b (#17510 ) * all2all * nan * remove that * less * has_local * only the nan change here * use nice getitem syntax for INDEX * work * remove * even simpler	2026-08-13 04:07:30 +09:00
RaineandGitHub	de04781b36	simplify equivalent const max (#17505 ) * add const max folds * add regression test * move	2026-08-12 08:39:29 -07:00
nimlgenandGitHub	4f106ebe87	hcq2: enqueue speed (#17504 )	2026-08-12 13:08:53 +03:00
qazalandGitHub	04c271ac41	simplify digitalocean_mi350x (#17502 ) * simplify digitalocean_mi350x * no hardcoded rocm path	2026-08-12 16:00:04 +09:00
qazalandGitHub	2e5a9a4121	no hardcoded device names in test_sliced_buffer_function (#17501 )	2026-08-12 15:19:25 +09:00
qazalandGitHub	e1013a6356	llama: create dataset cache by default in dev_beam (#17500 )	2026-08-12 15:03:39 +09:00
wozeparrotandGitHub	f891f5ffd0	gptoss: route lm_head thru asm_gemm (#17497 )	2026-08-11 18:32:50 -07:00
George HotzandGitHub	3686a1758f	mac/rdma imports lazy (#17496 ) * mac/rdma imports lazy * ish * fixes	2026-08-11 17:23:11 -07:00
George HotzandGitHub	4a253db9b4	minor cleanups to improve import speed (#17495 ) * minor cleanups to improve import speed * dumb	2026-08-11 16:06:25 -07:00
sirhcmandGitHub	479ffb0cda	remove Ops.SLICE (#17492 )	2026-08-11 18:50:04 -04:00
nimlgenandGitHub	2b5018e86a	hcq2: fix debug 2 info (#17491 ) * hcq2: fix debug 2 info * x * x * x	2026-08-12 00:25:44 +03:00
George HotzandGitHub	e11df72e0f	notes from digitalocean_mi350x (#17494 ) * notes from digitalocean_mi350x * cleanup * revert non-doc changes on digitalocean_mi350x branch	2026-08-11 13:22:11 -07:00
nimlgenandGitHub	a8c84ab34e	hcq2: enable all multitesnor tests (#17490 )	2026-08-11 17:47:33 +03:00
nimlgenandGitHub	ffef35c53e	hcq2: fix deps (#17481 ) * hcq2: proper unmap * hcq2: fix deps * x * x	2026-08-11 16:22:43 +03:00
nimlgenandGitHub	55e4f9d4f3	hcq2: proper unmap (#17489 )	2026-08-11 15:41:28 +03:00
sirhcmandGitHub	0c6a2c7dd6	slice is just shrink (#17483 )	2026-08-10 23:37:17 -04:00
RaineandGitHub	ad2fdeae69	move WMMA pms to codegen (#17485 ) * move wmma pms to codegen * lint tabs	2026-08-10 17:05:21 -07:00
RaineandGitHub	115bf9940f	add kwargs to group (#17484 )	2026-08-10 17:04:38 -07:00
geohot	a8d13380f9	simpler, but not inverse	2026-06-25 10:24:01 -07:00
geohot	afb0463f10	fix bitcast	2026-06-25 10:11:13 -07:00
geohot	947d6a3c69	enforce a cast for bitcast	2026-06-25 10:00:14 -07:00